Update RenderCUDA.py

TNTwise · May 12, 2024 · f45e8fe · f45e8fe
1 parent 2d626f7
commit f45e8fe
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/torch/RenderCUDA.py b/src/torch/RenderCUDA.py
@@ -375,11 +375,11 @@ def __init__(
         self.handleModel()
 
     def handleModel(self):
-        if "cuda" in self.method and "ncnn" not in self.method:
+        if "cuda1" in self.method and "ncnn" not in self.method:
             self.upscaleMethod = UpscaleCUDA(
                 self.originalWidth, self.originalHeight, self.model_path, self.half
             )
-        if "tensorrt" in self.method and "ncnn" not in self.method:
+        if "cuda" in self.method and "ncnn" not in self.method:
             self.upscaleMethod = UpscaleTensorRT(
                 width=self.originalWidth, height=self.originalHeight, modelPath=self.model_path, half=self.half
             )