Fixing synchronization issue when non-zero GPU is used as a device (#113)

Artyom17 · web-flow · commit b262949c2e16 · 2024-02-28T15:18:13.000-08:00
diff --git a/generate.py b/generate.py
@@ -15,7 +15,7 @@
 
 def device_sync(device):
     if "cuda" in device:
-        torch.cuda.synchronize()
+        torch.cuda.synchronize(device)
     elif "cpu" in device:
         pass
     else:
diff --git a/mixtral-moe/generate.py b/mixtral-moe/generate.py
@@ -15,7 +15,7 @@
 
 def device_sync(device):
     if "cuda" in device:
-        torch.cuda.synchronize()
+        torch.cuda.synchronize(device)
     elif "cpu" in device:
         pass
     else:
