Fix calling torch.mps.synchronize()

9 months ago · 5177079239
parent 9f3671114e
commit 5177079239
1 changed files with 1 additions and 1 deletions
--- a/src/petals/server/throughput.py
+++ b/src/petals/server/throughput.py
@ -233,7 +233,7 @@ def synchronize(device: torch.device):
    if device.type == "cuda":
        torch.cuda.synchronize(device)
    elif device.type == "mps":
-        torch.mps.synchronize(device)
+        torch.mps.synchronize()


 def get_device_name(device: torch.device) -> str: