Skip to content

GitLab

"ssh:/git@developer.sourcefind.cn:2222/OpenDAS/vllm_cscc.git" did not exist on "5aefd6ac3169b7b56023549cfa9614274d6e15f0"
Commit d9ef7ce7 authored by zhuwenwen's avatar zhuwenwen
Browse files

update rocm.py

parent 9e053941
Show whitespace changes
Inline Side-by-side
Showing with 4 additions and 3 deletions
vllm/platforms/rocm.py
View file @ d9ef7ce7
...@@ -225,8 +225,9 @@ class RocmPlatform(Platform): ...@@ -225,8 +225,9 @@ class RocmPlatform(Platform):
device: Optional[torch.types.Device] = None device: Optional[torch.types.Device] = None
) -> float: ) -> float:
torch.cuda.reset_peak_memory_stats(device) torch.cuda.reset_peak_memory_stats(device)
return torch.cuda.mem_get_info(device)[1] - torch.cuda.mem_get_info( # return torch.cuda.mem_get_info(device)[1] - torch.cuda.mem_get_info(
device)[0] # device)[0]
return torch.cuda.max_memory_allocated(device)
@classmethod @classmethod
def get_device_communicator_cls(cls) -> str: def get_device_communicator_cls(cls) -> str:
......
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