From 277f538d45e5a9309c006f03d48af1c2c37d6a14 Mon Sep 17 00:00:00 2001
From: CaRoLZhangxy <ucaszxy@outlook.com>
Date: Fri, 8 Mar 2024 08:31:57 +0000
Subject: [PATCH] fix type of typemap

---
 deepmd/pt/utils/dataset.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/deepmd/pt/utils/dataset.py b/deepmd/pt/utils/dataset.py
index 6bd92e8205..dbe4d92a0f 100644
--- 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: