[2 liner] TorchDataset.__len__ (#1268)

activeloopai · Oct 18, 2021 · 269833e · 269833e
1 parent a015ac2
commit 269833e
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 0 deletions.
diff --git a/hub/integrations/pytorch/pytorch.py b/hub/integrations/pytorch/pytorch.py
@@ -64,6 +64,8 @@ def __init__(
             if next_storage is not None:
                 next_storage.clear()
 
+            self.hub_dataset = dataset
+
             try:
                 self.cache = cache(
                     cache_storage=cache_storage,
@@ -80,6 +82,9 @@ def __init__(
                     "Underlying storage of the dataset in MemoryProvider which is not supported."
                 )
 
+        def __len__(self):
+            return len(self.hub_dataset)
+
         def __iter__(self):
             for value in self.cache.iterate_samples():
                 if value is not None:

diff --git a/hub/integrations/tests/test_pytorch.py b/hub/integrations/tests/test_pytorch.py
@@ -57,6 +57,8 @@ def test_pytorch_small(ds):
 
     dl = ds.pytorch(num_workers=2, batch_size=1)
 
+    assert len(dl.dataset) == 16
+
     for _ in range(2):
         for i, batch in enumerate(dl):
             np.testing.assert_array_equal(