Use inference mode in _MergedInferenceStep (#275)

1 year ago · fb2583b682
parent fd9400b392
commit fb2583b682
1 changed files with 1 additions and 0 deletions
--- a/src/petals/server/backend.py
+++ b/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,