fix(api): use training device when loading Real ESRGAN model (#67)

ssube · Jan 21, 2023 · 8c9c99e · 8c9c99e
1 parent 5286c4f
commit 8c9c99e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/api/onnx_web/convert.py b/api/onnx_web/convert.py
@@ -53,7 +53,7 @@
 model_path = environ.get('ONNX_WEB_MODEL_PATH',
                          path.join('..', 'models'))
 training_device = 'cuda' if torch.cuda.is_available() else 'cpu'
-map_location = None if torch.cuda.is_available() else torch.device('cpu')
+map_location = torch.device(training_device)
 
 @torch.no_grad()
 def convert_real_esrgan(name: str, url: str, scale: int, opset: int):