Merge pull request #116 from ndalton12/patch-1

Update loader.py to add DDP support
PhoenixDL · Jul 26, 2021 · ca0cf77 · ca0cf77
2 parents 7da711c + d437b09
commit ca0cf77
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/rising/loading/loader.py b/rising/loading/loader.py
@@ -189,6 +189,7 @@ def __init__(self, dataset: Union[Sequence, Dataset],
             gpu_transforms = Compose(to_gpu_trafo, gpu_transforms)
             gpu_transforms = gpu_transforms.to(device)
 
+        self.device = device
         self.sample_transforms = sample_transforms
         self.pseudo_batch_dim = pseudo_batch_dim and sample_transforms is not None
         self.gpu_transforms = gpu_transforms