瀏覽代碼

fix offloading lock

justheuristic 3 年之前
父節點
當前提交
db8607d571
共有 1 個文件被更改,包括 2 次插入2 次删除
  1. 2 2
      hivemind/optim/experimental/state_averager.py

+ 2 - 2
hivemind/optim/experimental/state_averager.py

@@ -545,8 +545,8 @@ class TrainingStateAverager(DecentralizedAverager):
             for local_param, loaded_param in zip(main_parameters_and_extras, loaded_parameters_and_extras):
                 local_param.copy_(loaded_param, non_blocking=True)
 
-            if self.offload_optimizer:
-                self._apply_optimizer_parameters_()
+        if self.offload_optimizer:
+            self._apply_optimizer_parameters_()
 
         self.local_epoch = metadata["epoch"]
         self._update_scheduler()