Update load weights

UTokyo-ICEPP · Jun 6, 2023 · fe09427 · fe09427
1 parent 44fe5b3
commit fe09427
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/multiml/task/pytorch/pytorch_base.py b/multiml/task/pytorch/pytorch_base.py
@@ -241,13 +241,14 @@ def load_model(self):
 
         # partial weights
         if ':' in model_path:
-            model_path, partial = model_path.split(':')
+            model_path, *partial = model_path.split(':')
             model_dict = self.ml.model.state_dict()
             state_dict = torch.load(model_path)
             new_state_dict = {}
             for key, value in state_dict.items():
-                if partial in key:
-                    new_state_dict[key] = value
+                for ipartial in partial:
+                    if ipartial in key:
+                        new_state_dict[key] = value
             model_dict.update(new_state_dict)
             self.ml.model.load_state_dict(model_dict)
         else: