fix jit for normal

iProzd · Sep 25, 2024 · 53adea0 · 53adea0
1 parent 78aae1f
commit 53adea0
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 0 deletions.
diff --git a/deepmd/pt/model/descriptor/repformers.py b/deepmd/pt/model/descriptor/repformers.py
@@ -492,6 +492,8 @@ def forward(
         if not self.direct_dist:
             g2, h2 = torch.split(dmatrix, [1, 3], dim=-1)
             if self.custom_radial:
+                assert self.radial_module is not None
+                assert self.cutoff_module is not None
                 rr = torch.linalg.norm(diff, dim=-1)
                 g2 = self.radial_module(rr) * self.cutoff_module(rr).unsqueeze(-1)
                 g2 = g2.view(nframes, nloc, nnei, -1)

diff --git a/deepmd/pt/model/descriptor/se_atten.py b/deepmd/pt/model/descriptor/se_atten.py
@@ -566,6 +566,8 @@ def forward(
             if not self.custom_radial:
                 ss = rr[:, :, :1]
             else:
+                assert self.radial_module is not None
+                assert self.cutoff_module is not None
                 dist = torch.linalg.norm(diff, dim=-1)
                 ss = self.radial_module(dist) * self.cutoff_module(dist).unsqueeze(-1)
                 ss = ss.view(nframes * nloc, nnei, -1)