Skip to content

Commit cf92544

Browse files
authored
Calls tensor.untyped_storage instead of tesnor.storage (#59)
1 parent 212e595 commit cf92544

File tree

1 file changed

+2
-2
lines changed

1 file changed

+2
-2
lines changed

pytorch_memlab/mem_reporter.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -92,7 +92,7 @@ def get_tensor_stat(tensor: torch.Tensor) -> List[Tuple[str, int, int, int]]:
9292

9393
numel = tensor.numel()
9494
element_size = tensor.element_size()
95-
fact_numel = tensor.storage().size()
95+
fact_numel = tensor.untyped_storage().size()
9696
fact_memory_size = fact_numel * element_size
9797
# since pytorch allocate at least 512 Bytes for any tensor, round
9898
# up to a multiple of 512
@@ -101,7 +101,7 @@ def get_tensor_stat(tensor: torch.Tensor) -> List[Tuple[str, int, int, int]]:
101101

102102
# tensor.storage should be the actual object related to memory
103103
# allocation
104-
data_ptr = tensor.storage().data_ptr()
104+
data_ptr = tensor.untyped_storage().data_ptr()
105105
if data_ptr in visited_data:
106106
name = '{}(->{})'.format(
107107
name,

0 commit comments

Comments
 (0)