Update test_utils.py

2024-03-20 11:59:24 +08:00
parent d2b0821990
commit bfea175a17
1 changed files with 2 additions and 1 deletions
@@ -156,7 +156,8 @@ def load_grl(generator_weight_PATH, scale=4):
            out_proj_type = "linear",
            conv_type = "1conv",
            upsampler = "nearest+conv",     # Change
-        ).cuda()
+        )
+        generator = generator.to(device)

    else:
        print("This weight is not supported")