Browse Source

Fix MPFuture failing outside inference mode (#521)

(cherry picked from commit 8f258b4b3688f671208bf323359cb967b25d640a)
Alexander Borzunov 2 năm trước cách đây
mục cha
commit
d9a986e443
1 tập tin đã thay đổi với 2 bổ sung1 xóa
  1. 2 1
      hivemind/utils/mpfuture.py

+ 2 - 1
hivemind/utils/mpfuture.py

@@ -127,7 +127,8 @@ class MPFuture(base.Future, Generic[ResultType]):
 
     @_state.setter
     def _state(self, new_state: State):
-        self._shared_state_code[...] = ALL_STATES.index(new_state)
+        with torch.inference_mode():
+            self._shared_state_code[...] = ALL_STATES.index(new_state)
         if self._state in TERMINAL_STATES and self._loop is not None and not self._aio_event.is_set():
             self._set_event_threadsafe()