Remove hard file offset reset in load() (#3695)

philippslang · soumith · commit c4b0db5079cc · 2017-11-17T15:21:37.000-05:00
* improved file offset logic

* load offset test

* whitespace

* needless exception handling

* test integer in binary
diff --git a/test/test_torch.py b/test/test_torch.py
@@ -8,6 +8,7 @@
 import tempfile
 import unittest
 import warnings
+import pickle
 from torch.utils.dlpack import from_dlpack, to_dlpack
 from itertools import product, combinations
 from common import TestCase, iter_indices, TEST_NUMPY, run_tests, download_file, skipIfNoLapack, \
@@ -4134,6 +4135,18 @@ def test_serialization(self):
             rootview = c[8]
             self.assertEqual(rootview.data_ptr(), c[0].data_ptr())
 
+    def test_serialization_offset(self):
+        a = torch.randn(5, 5)
+        i = 41
+        with tempfile.TemporaryFile() as f:
+            pickle.dump(i, f)
+            torch.save(a, f)
+            f.seek(0)
+            j = pickle.load(f)
+            b = torch.load(f)
+            self.assertTrue(torch.equal(a, b))
+            self.assertEqual(i, j)
+
     def test_half_tensor(self):
         x = torch.randn(5, 5).float()
         y = torch.randn(5, 5).float()
diff --git a/torch/serialization.py b/torch/serialization.py
@@ -389,13 +389,15 @@ def persistent_load(saved_id):
         else:
             raise RuntimeError("Unknown saved id type: %s" % saved_id[0])
 
-    # try the legacy loader first, which only works if f is a tarfile
-    try:
-        return legacy_load(f)
-    except tarfile.TarError:
-        pass
+    foffset = f.tell()
+    if foffset == 0:
+        # only if offset is zero we can attempt the legacy tar file loader
+        try:
+            return legacy_load(f)
+        except tarfile.TarError:
+            # if not a tarfile, reset file offset and proceed
+            f.seek(foffset)
 
-    f.seek(0)
     magic_number = pickle_module.load(f)
     if magic_number != MAGIC_NUMBER:
         raise RuntimeError("Invalid magic number; corrupt file?")
