Calculate memory usage in bytes for a tensor.
(tensor: torch.Tensor)
| 106 | input_tensors += [value for value in kwargs.values() if isinstance(value, torch.Tensor)] |
| 107 | |
| 108 | def calculate_memory(tensor: torch.Tensor): |
| 109 | """Calculate memory usage in bytes for a tensor.""" |
| 110 | return tensor.numel() * tensor.element_size() |
| 111 | |
| 112 | input_memory = sum(calculate_memory(t) for t in input_tensors) |
| 113 |