Merge pull request #2571 from activeloopai/batch_sampler

added torch batch_sampler

deeplake/enterprise/dataloader.py

@@ -25,10 +25,13 @@
 try:
     from torch.utils.data.dataloader import DataLoader, _InfiniteConstantSampler
     from torch.utils.data.distributed import DistributedSampler
+    from torch.utils.data import BatchSampler
+
 except ImportError:
     DataLoader = object  # type: ignore
     _InfiniteConstantSampler = None  # type: ignore
     DistributedSampler = None  # type: ignore
+    BatchSampler = None  # type: ignore
 
 import numpy as np
 
@@ -196,7 +199,7 @@ def sampler(self):
 
     @property
     def batch_sampler(self):
-        return DistributedSampler(self.dataset) if self._distributed else None
+        return BatchSampler(self.sampler, self.batch_size, self.drop_last) if BatchSampler else None
 
     @property
     def generator(self):

deeplake/enterprise/test_pytorch.py

@@ -722,6 +722,17 @@ def test_pytorch_error_handling(local_auth_ds):
         pass
 
 
+@requires_libdeeplake
+@requires_torch
+def test_batch_sampler_attribute(local_auth_ds):
+    ld = local_auth_ds.dataloader().pytorch()
+
+    from torch.utils.data import BatchSampler
+
+    assert isinstance(ld.batch_sampler, BatchSampler)
+    assert ld.batch_sampler.sampler is not None
+
+
 @requires_libdeeplake
 @requires_torch
 @pytest.mark.slow