fix(api): handle more SwinIR models

ssube · Jun 17, 2023 · 1506f51 · 1506f51
1 parent 12e489b
commit 1506f51
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/api/onnx_web/convert/upscaling/swinir.py b/api/onnx_web/convert/upscaling/swinir.py
@@ -86,8 +86,10 @@ def convert_upscaling_swinir(
     torch_model = torch.load(source, map_location=conversion.map_location)
     if "params_ema" in torch_model:
         model.load_state_dict(torch_model["params_ema"], strict=False)
-    else:
+    elif "params" in torch_model:
         model.load_state_dict(torch_model["params"], strict=False)
+    else:
+        model.load_state_dict(torch_model, strict=False)
 
     model.to(conversion.training_device).train(False)
     model.eval()