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