fix ut

Author: caic99

deepmd/pt/train/training.py

@@ -173,9 +173,7 @@ def cycle_iterator(iterable: Iterable):
             Any: The next item from the iterable, cycling back to the beginning when the end is reached.
             """
             while True:
-                with torch.device("cpu"):
-                    it = iter(iterable)
-                yield from it
+                yield from iterable
 
         def get_data_loader(_training_data, _validation_data, _training_params):
             def get_dataloader_and_iter(_data, _params):
@@ -194,7 +192,6 @@ def get_dataloader_and_iter(_data, _params):
                     drop_last=False,
                     collate_fn=lambda batch: batch,  # prevent extra conversion
                     pin_memory=True,
-                    pin_memory_device=str(object=DEVICE),
                 )
                 _data_iter = cycle_iterator(_dataloader)
                 return _dataloader, _data_iter

deepmd/pt/utils/dataloader.py

@@ -267,7 +267,16 @@ def get_weighted_sampler(training_data, prob_style, sys_prob=False):
     # training_data.total_batch is the size of one epoch, you can increase it to avoid too many  rebuilding of iterators
     len_sampler = training_data.total_batch * max(env.NUM_WORKERS, 1)
     with torch.device("cpu"):
-        sampler = WeightedRandomSampler(probs, len_sampler, replacement=True)
+        sampler = WeightedRandomSampler(
+            probs,
+            len_sampler,
+            replacement=True,
+            generator=torch.Generator(),
+            # If we are not setting the generator here, the random state will be initialized in the beginning of each new epoch.
+            # This operation involves creating a new tensor for seeding on the default device,
+            # while unit tests requires specifying the device explicitly.
+            # https://github.com/pytorch/pytorch/blob/134179474539648ba7dee1317959529fbd0e7f89/torch/utils/data/sampler.py#L170
+        )
     return sampler