speed up onnx output

TNTwise · May 23, 2024 · ad8a89c · ad8a89c
1 parent 7c65e31
commit ad8a89c
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/torch/UpscaleImageTensorRT.py b/src/torch/UpscaleImageTensorRT.py
@@ -94,7 +94,7 @@ def pytorchExportToONNX(self):  # Loads model via spandrel, and exports to onnx
         state_dict = model.state_dict()
         model.eval().cuda()
         model.load_state_dict(state_dict, strict=True)
-        input = torch.rand(1, 3, 256, 256).cuda()
+        input = torch.rand(1, 3, 20, 20).cuda()
         if self.half:
             try:
                 model.half()