apextoaster
/
onnx-web
1
0
Fork 0
Code Actions

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

This commit is contained in:
Sean Sube 2023-01-21 17:55:55 -06:00
parent 5286c4f596
commit 8c9c99eeb5
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
api/onnx_web/convert.py
View File

@ -53,7 +53,7 @@ extra_models: Models = {
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):