diff --git a/compressai/models/video/google.py b/compressai/models/video/google.py
index b978af8d..1dee5044 100644
--- a/compressai/models/video/google.py
+++ b/compressai/models/video/google.py
@@ -359,7 +359,7 @@ def warp_volume(self, volume, flow, scale_field, padding_mode: str = "border"):
 
         N, C, _, H, W = volume.size()
 
-        with amp.autocast(device_type=volume.device.type, enabled=False):
+        with amp.autocast(volume.device.type, enabled=False):
             grid = meshgrid2d(N, C, H, W, volume.device)
             update_grid = grid + flow.permute(0, 2, 3, 1).float()
             update_scale = scale_field.permute(0, 2, 3, 1).float()
diff --git a/compressai/utils/video/eval_model/__main__.py b/compressai/utils/video/eval_model/__main__.py
index 768ae197..295eda5a 100644
--- a/compressai/utils/video/eval_model/__main__.py
+++ b/compressai/utils/video/eval_model/__main__.py
@@ -371,7 +371,7 @@ def run_inference(
             continue
 
         with amp.autocast(
-            device_type=next(net.parameters()).device.type, enabled=args["half"]
+            next(net.parameters()).device.type, enabled=args["half"]
         ):
             with torch.no_grad():
                 if entropy_estimation: