fix: try fix multitask

deepmodeling · Apr 2, 2024 · bd89e3a · bd89e3a
1 parent 194de07
commit bd89e3a
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/deepmd/pt/model/task/fitting.py b/deepmd/pt/model/task/fitting.py
@@ -326,7 +326,8 @@ def get_sel_type(self) -> List[int]:
     def __setitem__(self, key, value):
         if key in ["bias_atom_e"]:
             if isinstance(value, dict):
-                value = {k: v.view([self.ntypes, -1]) for k, v in value.items()}
+                #TODO: need to update to handle multiple properties
+                value = list(value.values())[0].view([self.ntypes, self._net_out_dim()])
             else:
                 value = value.view([self.ntypes, self._net_out_dim()])
             self.bias_atom_e = value