From bdf659aad5a112c4ecd9bbb7198fc870ec2b169c Mon Sep 17 00:00:00 2001 From: Anthony Mercurio Date: Tue, 14 Nov 2023 17:43:51 -0700 Subject: [PATCH] fix: Correct tensor size calculation in TensorSerializer --- tensorizer/serialization.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorizer/serialization.py b/tensorizer/serialization.py index b2edffb..7fb7d3e 100644 --- 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)