Explorar o código

Use inference mode in _MergedInferenceStep (#275)

justheuristic %!s(int64=2) %!d(string=hai) anos
pai
achega
fb2583b682
Modificáronse 1 ficheiros con 1 adicións e 0 borrados
  1. 1 0
      src/petals/server/backend.py

+ 1 - 0
src/petals/server/backend.py

@@ -159,6 +159,7 @@ class _MergedInferenceStep:
     def __init__(self, backends: Dict[ExpertUID, TransformerBackend]):
         self.backends = backends
 
+    @torch.inference_mode()
     def __call__(
         self,
         hidden_states: torch.Tensor,