[train_gpt2.py] synchronize based on device

czzhao-sjtu · Apr 11, 2024 · 4542f89 · 4542f89
1 parent a08c11b
commit 4542f89
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/train_gpt2.py b/train_gpt2.py
@@ -386,7 +386,10 @@ def get_batch():
                 write_model(model, "gpt2_124M.bin")
                 write_state(model, x, y, logits, loss, "gpt2_124M_debug_state.bin")
             optimizer.step()
-        torch.cuda.synchronize()
+        if device == "mps":
+            torch.mps.synchronize()
+        elif device == "cuda":
+            torch.cuda.synchronize()
         t1 = time.time()
         print(f"iteration {i}, loss: {loss.item()}, time: {(t1-t0)*1000:.3f}ms")