fix missing sync

pytorch · Feb 1, 2025 · 471411d · 471411d
1 parent e65cfac
commit 471411d
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/torchao/prototype/low_bit_optim/cpu_offload.py b/torchao/prototype/low_bit_optim/cpu_offload.py
@@ -107,6 +107,8 @@ def step(self, closure=None):
             with getattr(torch, self.device).stream(self.stream):
                 p_device.copy_(p_host, non_blocking=True)
 
+        # make sure param H2D finishes before the next forward pass
+        self.stream.synchronize()
         self.queue.clear()
         return loss