change | to Union

deepmodeling · Dec 13, 2024 · 4c2033f · 4c2033f
1 parent 0608e17
commit 4c2033f
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 4 deletions.
diff --git a/deepmd/dpmodel/output_def.py b/deepmd/dpmodel/output_def.py
@@ -197,7 +197,7 @@ def __init__(
         r_hessian: bool = False,
         magnetic: bool = False,
         intensive: bool = False,
-        sub_var_name: list = None,
+        sub_var_name: list[str] = None,
     ) -> None:
         self.name = name
         self.shape = list(shape)

diff --git a/deepmd/pt/loss/property.py b/deepmd/pt/loss/property.py
@@ -3,6 +3,7 @@
 
 import torch
 import torch.nn.functional as F
+from typing import Union
 
 from deepmd.pt.loss.loss import (
     TaskLoss,
@@ -24,8 +25,8 @@ def __init__(
         loss_func: str = "smooth_mae",
         metric: list = ["mae"],
         beta: float = 1.00,
-        property_name: str | list = "property",
-        property_dim: int | list = 1,
+        property_name: Union[str, list] = "property",
+        property_dim: Union[int, list] = 1,
         **kwargs,
     ) -> None:
         r"""Construct a layer to compute loss on property.

diff --git a/deepmd/pt/model/task/property.py b/deepmd/pt/model/task/property.py
@@ -82,7 +82,7 @@ def __init__(
         bias_atom_p: Optional[torch.Tensor] = None,
         intensive: bool = False,
         bias_method: str = "normal",
-        property_name: str | list = "property",
+        property_name: Union[str, list] = "property",
         resnet_dt: bool = True,
         numb_fparam: int = 0,
         numb_aparam: int = 0,