Update model.py

601b568f · Maksim Skobeev · GitHub · a9a671cb · 601b568f
Unverified Commit 601b568f authored Dec 22, 2023 by Maksim Skobeev Committed by GitHub Dec 22, 2023
Show whitespace changes
Inline Side-by-side

Showing with 2 additions and 1 deletion

RealESRGAN/model.py RealESRGAN/model.py +2 -1

No files found.
--- a/RealESRGAN/model.py
+++ b/RealESRGAN/model.py
@@ -56,9 +56,10 @@ class RealESRGAN:
        self.model.eval()
        self.model.to(self.device)
        
-    @torch.cuda.amp.autocast()
+    # @torch.cuda.amp.autocast()
    def predict(self, lr_image, batch_size=4, patches_size=192,
                padding=24, pad_size=15):
+        torch.autocast(device_type=self.device.type)
        scale = self.scale
        device = self.device
        lr_image = np.array(lr_image)