update seed setup

XinhaoMei · Jan 11, 2022 · 8cc1abc · 8cc1abc
1 parent 31b8d6e
commit 8cc1abc
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/tools/utils.py b/tools/utils.py
@@ -17,7 +17,9 @@
 def setup_seed(seed):
 
     torch.manual_seed(seed)
-    torch.cuda.manual_seed_all(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)
     np.random.seed(seed)
     random.seed(seed)
     torch.backends.cudnn.deterministic = True