소스 검색

Fix MPFuture failing outside inference mode (#521)

Alexander Borzunov 2 년 전
부모
커밋
8f258b4b36
1개의 변경된 파일2개의 추가작업 그리고 1개의 파일을 삭제
  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()