add Optional

Signed-off-by: Jinzhe Zeng <[email protected]>
deepmodeling · Jan 28, 2024 · b3f0b8d · b3f0b8d
1 parent a3671d7
commit b3f0b8d
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/deepmd/utils/random.py b/deepmd/utils/random.py
@@ -12,7 +12,7 @@
 
 def choice(
     a: Union[np.ndarray, int],
-    size: Union[int, Tuple[int, ...]] = None,
+    size: Optional[Union[int, Tuple[int, ...]]] = None,
     replace: bool = True,
     p: Optional[np.ndarray] = None,
 ):