aboutsummaryrefslogtreecommitdiff
path: root/modules/swinir_model.py
diff options
context:
space:
mode:
authorAUTOMATIC1111 <16777216c@gmail.com>2022-12-03 09:58:08 +0300
committerGitHub <noreply@github.com>2022-12-03 09:58:08 +0300
commita2feaa95fc0c4c94131eb75b5b1bc0eaa1696551 (patch)
treea7be0e8b6849aae24f7f6c7879f7ddf43d118425 /modules/swinir_model.py
parentc7af672186ec09a514f0e78aa21155264e56c130 (diff)
parent0fddb4a1c06a6e2122add7eee3b001a6d473baee (diff)
Merge pull request #5194 from brkirch/autocast-and-mps-randn-fixes
Use devices.autocast() and fix MPS randn issues
Diffstat (limited to 'modules/swinir_model.py')
-rw-r--r--modules/swinir_model.py6
1 files changed, 1 insertions, 5 deletions
diff --git a/modules/swinir_model.py b/modules/swinir_model.py
index facd262d..483eabd4 100644
--- a/modules/swinir_model.py
+++ b/modules/swinir_model.py
@@ -13,10 +13,6 @@ from modules.swinir_model_arch import SwinIR as net
from modules.swinir_model_arch_v2 import Swin2SR as net2
from modules.upscaler import Upscaler, UpscalerData
-precision_scope = (
- torch.autocast if cmd_opts.precision == "autocast" else contextlib.nullcontext
-)
-
class UpscalerSwinIR(Upscaler):
def __init__(self, dirname):
@@ -112,7 +108,7 @@ def upscale(
img = np.moveaxis(img, 2, 0) / 255
img = torch.from_numpy(img).float()
img = img.unsqueeze(0).to(devices.device_swinir)
- with torch.no_grad(), precision_scope("cuda"):
+ with torch.no_grad(), devices.autocast():
_, _, h_old, w_old = img.size()
h_pad = (h_old // window_size + 1) * window_size - h_old
w_pad = (w_old // window_size + 1) * window_size - w_old