Update torch.cuda.amp to torch.amp in pytorch.py (#945)

update pytorch.py to get rid of torch.cuda.amp deprecated warning.
explosion · Sep 30, 2024 · 12e308c · 12e308c
1 parent f2166b4
commit 12e308c
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/thinc/shims/pytorch.py b/thinc/shims/pytorch.py
@@ -111,7 +111,7 @@ def predict(self, inputs: ArgsKwargs) -> Any:
         """
         self._model.eval()
         with torch.no_grad():
-            with torch.cuda.amp.autocast(self._mixed_precision):
+            with torch.amp.autocast("cuda", self._mixed_precision):
                 outputs = self._model(*inputs.args, **inputs.kwargs)
         self._model.train()
         return outputs
@@ -125,7 +125,7 @@ def begin_update(self, inputs: ArgsKwargs):
         self._model.train()
 
         # Note: mixed-precision autocast must not be applied to backprop.
-        with torch.cuda.amp.autocast(self._mixed_precision):
+        with torch.amp.autocast("cuda", self._mixed_precision):
             output = self._model(*inputs.args, **inputs.kwargs)
 
         def backprop(grads):