Merge pull request #457 from activeloopai/fix/to_pytorch_flush

Hotfix for to_pytorch
activeloopai · Jan 15, 2021 · 2a9a7e1 · 2a9a7e1
2 parents a394e9b + 6a0486d
commit 2a9a7e1
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 4 deletions.
diff --git a/hub/api/dataset.py b/hub/api/dataset.py
@@ -519,13 +519,14 @@ def to_pytorch(
         num_samples: int, optional
             The number of samples required of the dataset that needs to be converted
         """
-        if "torch" not in sys.modules:
+        try:
+            import torch
+        except ModuleNotFoundError:
             raise ModuleNotInstalledException("torch")
-        import torch
 
         global torch
-
-        self.flush()  # FIXME Without this some tests in test_converters.py fails, not clear why
+        if "r" not in self.mode:
+            self.flush()  # FIXME Without this some tests in test_converters.py fails, not clear why
         return TorchDataset(
             self,
             transform,

diff --git a/hub/api/tests/test_converters.py b/hub/api/tests/test_converters.py
@@ -334,6 +334,12 @@ def test_to_from_pytorch():
         assert (res_ds["label", "d", "e", i].numpy() == i * np.ones((5, 3))).all()
 
 
+@pytest.mark.skipif(not pytorch_loaded(), reason="requires pytorch to be loaded")
+def test_to_pytorch_bug():
+    ds = hub.Dataset("activeloop/mnist", mode="r")
+    data = ds.to_pytorch()
+
+
 if __name__ == "__main__":
     with Timer("Test Converters"):
         with Timer("from MNIST"):