Update utils.py

deepmodeling · May 14, 2024 · f3da347 · f3da347
1 parent 29473c5
commit f3da347
Showing 1 changed file with 1 addition and 3 deletions.
diff --git a/deepmd/pt/utils/utils.py b/deepmd/pt/utils/utils.py
@@ -113,8 +113,6 @@ def dict_to_device(sample_dict):
 
 def get_generator(seed: Optional[int] = None) -> Optional[torch.Generator]:
     if seed is not None:
-        generator = torch.Generator(device=DEVICE)
-        generator.manual_seed(seed)
-        return generator
+        return torch.manual_seed(seed)
     else:
         return None