fix: wrap _import_static_state in inference_mode to fix resume on Blackwell (#21035)

2026-04-08 02:03:39 -07:00
parent 79c82c5c42
commit 413913763f
1 changed files with 4 additions and 3 deletions
@@ -220,6 +220,7 @@ def _export_static_state(model):


 def _import_static_state(model, static_params):
-    self_named_buffers = dict(model.named_buffers())
-    for name, tensor in static_params["buffers"]:
-        self_named_buffers[name][...] = tensor
+    with torch.inference_mode():
+        self_named_buffers = dict(model.named_buffers())
+        for name, tensor in static_params["buffers"]:
+            self_named_buffers[name][...] = tensor