yujunjun/haystack
1
0
Fork 0
mirror of https://github.com/deepset-ai/haystack.git synced 2026-01-06 03:57:19 +00:00
Code Issues Packages Projects Releases Wiki Activity

Check cuda availability before calling (#4174)

Browse Source 
Co-authored-by: Massimiliano Pippi <mpippi@gmail.com>
This commit is contained in:
abwiersma 2023-02-20 17:37:56 +01:00 committed by GitHub
parent 18e7b8399b
commit 7aae4293d7
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
haystack/modeling/training/base.py
View File

@ -505,7 +505,7 @@ class Trainer:
"scheduler_state": self.lr_schedule.state_dict(),
"numpy_rng_state": numpy.random.get_state(),
"rng_state": torch.get_rng_state(),
"cuda_rng_state": torch.cuda.get_rng_state(),
"cuda_rng_state": torch.cuda.get_rng_state() if torch.cuda.is_available() else None,
},
checkpoint_path / "trainer",
pickle_module=dill,