Update inference.py

Author: iProzd

deepmd/pt/infer/inference.py

@@ -40,6 +40,8 @@ def __init__(
         model_params = state_dict["_extra_state"]["model_params"]
         self.multi_task = "model_dict" in model_params
         if self.multi_task:
+            if head is None and "Default" in model_params["model_dict"]:
+                head = "Default"
             assert head is not None, "Head must be specified in multitask mode!"
             self.head = head
             assert head in model_params["model_dict"], (