Fix setup_seed function to handle non-list seed

njzjz · Dec 19, 2024 · 124f432 · 124f432
1 parent b95718a
commit 124f432
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/deepmd/pt/utils/dataloader.py b/deepmd/pt/utils/dataloader.py
@@ -52,6 +52,8 @@
 def setup_seed(seed) -> None:
     if isinstance(seed, (list, tuple)):
         mixed_seed = mix_entropy(seed)
+    else:
+        mixed_seed = seed
     torch.manual_seed(mixed_seed)
     torch.cuda.manual_seed_all(mixed_seed)
     torch.backends.cudnn.deterministic = True