fix(pt): Specify type for type_embd_data in DescrptBlockSeAtten

OutisLi · OutisLi · commit bffe4869ace9 · 2025-11-20T15:19:16.000+08:00
diff --git a/deepmd/pt/model/descriptor/se_atten.py b/deepmd/pt/model/descriptor/se_atten.py
@@ -285,7 +285,7 @@ def __init__(
             [nn.Parameter(torch.zeros(0, dtype=self.prec, device=env.DEVICE))]
         )
         # For type embedding compression
-        self.type_embd_data = None
+        self.type_embd_data: Optional[torch.Tensor] = None
 
     def get_rcut(self) -> float:
         """Returns the cut-off radius."""

Original file line number	Diff line number	Diff line change
`@@ -285,7 +285,7 @@ def __init__(`
`285`	`285`	`[nn.Parameter(torch.zeros(0, dtype=self.prec, device=env.DEVICE))]`
`286`	`286`	`)`
`287`	`287`	`# For type embedding compression`
`288`		`- self.type_embd_data = None`
	`288`	`+ self.type_embd_data: Optional[torch.Tensor] = None`
`289`	`289`
`290`	`290`	`def get_rcut(self) -> float:`
`291`	`291`	`"""Returns the cut-off radius."""`