From bba5e2179dc5161734f33bda6069f40b3ce048f7 Mon Sep 17 00:00:00 2001
From: Falko Habel <falko.habel@gmx.de>
Date: Mon, 17 Feb 2025 16:26:54 +0100
Subject: [PATCH] setted up correct amp

---
 src/aiunn/inference.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

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