up

IPL-UV · Oct 24, 2024 · a3194c0 · a3194c0
1 parent 3ec2b17
commit a3194c0
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/supers2/setup.py b/supers2/setup.py
@@ -271,7 +271,7 @@ def load_srx4_model(
     model_params["upscale"] = 4
 
     # Load the model
-    SRX4 = load_model(model_name, model_params, kwargs)
+    SRX4 = load_model(model_name, model_params, device=device, **kwargs)
     SRX4.load_state_dict(weights_data)
     SRX4.eval()
     for param in SRX4.parameters():