Fix EnergyFittingNetDirect

iProzd · Feb 28, 2024 · ae27607 · ae27607
1 parent 08e18fe
commit ae27607
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/deepmd/pt/model/model/dp_model.py b/deepmd/pt/model/model/dp_model.py
@@ -10,6 +10,7 @@
 )
 from deepmd.pt.model.task.ener import (
     EnergyFittingNet,
+    EnergyFittingNetDirect,
 )
 from deepmd.pt.model.task.polarizability import (
     PolarFittingNet,
@@ -36,7 +37,9 @@ def __new__(cls, descriptor, fitting, *args, **kwargs):
         # according to the fitting network to decide the type of the model
         if cls is DPModel:
             # map fitting to model
-            if isinstance(fitting, EnergyFittingNet):
+            if isinstance(fitting, EnergyFittingNet) or isinstance(
+                fitting, EnergyFittingNetDirect
+            ):
                 cls = EnergyModel
             elif isinstance(fitting, DipoleFittingNet):
                 cls = DipoleModel