fix type of typemap

deepmodeling · Mar 8, 2024 · 277f538 · 277f538
1 parent 8ae44f5
commit 277f538
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/deepmd/pt/utils/dataset.py b/deepmd/pt/utils/dataset.py
@@ -3,6 +3,7 @@
 
 from typing import (
     List,
+    Optional,
 )
 
 from torch.utils.data import (
@@ -16,7 +17,7 @@
 
 
 class DeepmdDataSetForLoader(Dataset):
-    def __init__(self, system: str, type_map: str):
+    def __init__(self, system: str, type_map: Optional[List[str]] = None):
         """Construct DeePMD-style dataset containing frames cross different systems.
 
         Args: