fix: Correct tensor size calculation in TensorSerializer

coreweave · Nov 15, 2023 · bdf659a · bdf659a
1 parent b38d959
commit bdf659a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tensorizer/serialization.py b/tensorizer/serialization.py
@@ -2557,7 +2557,7 @@ def _bulk_write(self, tensors: Iterable[_WriteSpec]):
         fallocate = getattr(os, "posix_fallocate", None)
         if fallocate and self._fd:
             size = sum(len(t.name) for t in tensors)
-            size += sum(t.tensor.untyped_storage().size() for t in tensors)
+            size += sum(t.tensor.element_size() * t.tensor.nelement() for t in tensors)
             # Rough underestimate of header size
             header_min_size = 24
             size += header_min_size * len(tensors)