Fix model loading (#915)

* Fix model loading * remove logging
chaiNNer-org · Sep 7, 2022 · df1e1a9 · df1e1a9
1 parent fed3a2f
commit df1e1a9
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/backend/src/nodes/utils/pytorch_model_loading.py b/backend/src/nodes/utils/pytorch_model_loading.py
@@ -15,11 +15,13 @@ def load_state_dict(state_dict) -> PyTorchModel:
 
     if "params_ema" in state_dict_keys:
         state_dict = state_dict["params_ema"]
-    if "params-ema" in state_dict_keys:
+    elif "params-ema" in state_dict_keys:
         state_dict = state_dict["params-ema"]
-    if "params" in state_dict_keys:
+    elif "params" in state_dict_keys:
         state_dict = state_dict["params"]
 
+    state_dict_keys = list(state_dict.keys())
+
     # SRVGGNet Real-ESRGAN (v2)
     if "body.0.weight" in state_dict_keys and "body.1.weight" in state_dict_keys:
         model = RealESRGANv2(state_dict)