Fix race condition that can corrupt grads under certain conditions.

2026-08-07 14:52:47 -06:00
parent 817f3dcbcb
commit f4e9130547
3 changed files with 53 additions and 15 deletions
--- a/toolkit/memory_management/init.py
+++ b/toolkit/memory_management/init.py
@@ -1 +1,2 @@
-from .manager import MemoryManager
+from .manager import MemoryManager
+from .manager_modules import sync_grad_transfers