keep old behaviors

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
deepmodeling · Mar 21, 2024 · 2a61e3c · 2a61e3c
1 parent 89e640c
commit 2a61e3c
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 0 deletions.
diff --git a/deepmd/tf/fit/dipole.py b/deepmd/tf/fit/dipole.py
@@ -224,6 +224,8 @@ def build(
             nframes = tf.shape(inputs)[0]
 
         if self.mixed_types or type_embedding is not None:
+            # keep old behavior
+            self.mixed_types = True
             nloc_mask = tf.reshape(
                 tf.tile(tf.repeat(self.sel_mask, natoms[2:]), [nframes]), [nframes, -1]
             )

diff --git a/deepmd/tf/fit/dos.py b/deepmd/tf/fit/dos.py
@@ -509,6 +509,8 @@ def build(
             tf.slice(atype_nall, [0, 0], [-1, natoms[0]]), [-1]
         )  ## lammps will make error
         if type_embedding is not None:
+            # keep old behavior
+            self.mixed_types = True
             atype_embed = tf.nn.embedding_lookup(type_embedding, self.atype_nloc)
         else:
             atype_embed = None

diff --git a/deepmd/tf/fit/ener.py b/deepmd/tf/fit/ener.py
@@ -635,6 +635,8 @@ def build(
         ):
             type_embedding = nvnmd_cfg.map["t_ebd"]
         if type_embedding is not None:
+            # keep old behavior
+            self.mixed_types = True
             atype_embed = tf.nn.embedding_lookup(type_embedding, self.atype_nloc)
         else:
             atype_embed = None

diff --git a/deepmd/tf/fit/polar.py b/deepmd/tf/fit/polar.py
@@ -403,6 +403,8 @@ def build(
             nframes = tf.shape(inputs)[0]
 
         if self.mixed_types or type_embedding is not None:
+            # keep old behavior
+            self.mixed_types = True
             # nframes x nloc
             nloc_mask = tf.reshape(
                 tf.tile(tf.repeat(self.sel_mask, natoms[2:]), [nframes]), [nframes, -1]