From 9086dfc73f25a7952d5c35e6eb0b2a3d6defbaaf Mon Sep 17 00:00:00 2001
From: Andre Anjos <andre.dos.anjos@gmail.com>
Date: Fri, 21 Jul 2023 00:57:51 +0200
Subject: [PATCH] [data.datamodule] Adjust sample storage size calculation
---
src/ptbench/data/datamodule.py | 3 ++-
1 file changed, 2 insertions(+), 1 deletion(-)
diff --git a/src/ptbench/data/datamodule.py b/src/ptbench/data/datamodule.py
index 4c924763..bb2dcdda 100644
--- a/src/ptbench/data/datamodule.py
+++ b/src/ptbench/data/datamodule.py
@@ -49,7 +49,8 @@ def _sample_size_bytes(s: Sample) -> int:
"""Returns a tensor size in bytes."""
return int(t.element_size() * torch.prod(torch.tensor(t.shape)))
- size = int(s[0].element_size() * torch.prod(torch.tensor(s[0].shape)))
+ size = sys.getsizeof(s[0]) # tensor metadata
+ size += int(s[0].element_size() * torch.prod(torch.tensor(s[0].shape)))
size += sys.getsizeof(s[1])
# check each element - if it is a tensor, then adds its total space in
--
GitLab