diff --git a/src/aiunn/inference.py b/src/aiunn/inference.py
index b364191..182c8da 100644
--- a/src/aiunn/inference.py
+++ b/src/aiunn/inference.py
@@ -124,7 +124,7 @@ class ImageUpscaler:
         
         # Perform inference with the model
         with torch.no_grad():
-            with torch.cuda.amp.autocast(device_type="cuda"):
+            with torch.amp.autocast(device_type="cuda"):
                 output_tensor = self.model(preprocessed_image)
         
         # Postprocess and return the upscaled image