Update mlp.py

Author: iProzd

deepmd/pt/model/network/mlp.py

@@ -521,7 +521,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
             P_list.append(Plp1)
 
         P = torch.concat(P_list, dim=-1)  # (..., L+1)
-        return P * self.norm
+        return P * self.norm.type(x.dtype)
 
 
 def find_normalization(name: str, dim: int | None = None) -> nn.Module | None: