Unverified Commit 7539bbc6 authored by Gregory Shtrasberg's avatar Gregory Shtrasberg Committed by GitHub
Browse files

[ROCm] Using a more precise memory profiling (#12624)


Signed-off-by: default avatarGregory Shtrasberg <Gregory.Shtrasberg@amd.com>
parent 9cf47594
......@@ -169,4 +169,5 @@ class RocmPlatform(Platform):
device: Optional[torch.types.Device] = None
) -> float:
torch.cuda.reset_peak_memory_stats(device)
return torch.cuda.max_memory_allocated(device)
return torch.cuda.mem_get_info(device)[1] - torch.cuda.mem_get_info(
device)[0]
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment