Răsfoiți Sursa

Fix MPFuture failing outside inference mode (#521)

Alexander Borzunov 2 ani în urmă
părinte
comite
8f258b4b36
1 a modificat fișierele cu 2 adăugiri și 1 ștergeri
  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()