Have __sizeof__ account for size of stored elements (#3821)

lantiga · ezyang · commit 9b31280ccf60 · 2017-11-27T11:22:34.000-05:00
* Have __sizeof__ account for size of stored elements

* Conform to sizeof specification
diff --git a/test/test_torch.py b/test/test_torch.py
@@ -4349,6 +4349,19 @@ def test_print(self):
         x.__repr__()
         str(x),
 
+    def test_sizeof(self):
+        sizeof_empty = torch.randn(0).storage().__sizeof__()
+        sizeof_10 = torch.randn(10).storage().__sizeof__()
+        sizeof_100 = torch.randn(100).storage().__sizeof__()
+        self.assertEqual((sizeof_100 - sizeof_empty) // (sizeof_10 - sizeof_empty), 10)
+        self.assertEqual((sizeof_100 - sizeof_empty) % (sizeof_10 - sizeof_empty), 0)
+
+        sizeof_empty = torch.randn(0).type(torch.ByteTensor).storage().__sizeof__()
+        sizeof_10 = torch.randn(10).type(torch.ByteTensor).storage().__sizeof__()
+        sizeof_100 = torch.randn(100).type(torch.ByteTensor).storage().__sizeof__()
+        self.assertEqual((sizeof_100 - sizeof_empty) // (sizeof_10 - sizeof_empty), 10)
+        self.assertEqual((sizeof_100 - sizeof_empty) % (sizeof_10 - sizeof_empty), 0)
+
     def test_unsqueeze(self):
         x = torch.randn(2, 3, 4)
         y = x.unsqueeze(1)
diff --git a/torch/storage.py b/torch/storage.py
@@ -30,6 +30,9 @@ def __deepcopy__(self, memo):
     def __reduce__(self):
         return type(self), (self.tolist(),)
 
+    def __sizeof__(self):
+        return super(_StorageBase, self).__sizeof__() + self.element_size() * self.size()
+
     def clone(self):
         """Returns a copy of this storage"""
         return type(self)(self.size()).copy_(self)
