address comments

apax-hub · Nov 28, 2024 · 882300e · 882300e
1 parent 46fb4da
commit 882300e
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/apax/nodes/md.py b/apax/nodes/md.py
@@ -26,7 +26,7 @@ class ApaxJaxMD(zntrack.Node):
         index of the configuration from the data list to use
     model: ApaxModel
         model to use for the simulation
-    repeat: float
+    repeat: None|int|tuple[int, int, int]
         number of repeats
     config: str
         path to the MD simulation parameter file
@@ -36,7 +36,7 @@ class ApaxJaxMD(zntrack.Node):
     data_id: int = zntrack.params(-1)
 
     model: ApaxBase = zntrack.deps()
-    repeat: typing.Optional[bool] = zntrack.params(None)
+    repeat: None|int|tuple[int, int, int] = zntrack.params(None)
 
     config: str = zntrack.params_path(None)
 

diff --git a/apax/nodes/model.py b/apax/nodes/model.py
@@ -41,7 +41,7 @@ class Apax(ApaxBase):
         verbosity of logging during training
     """
 
-    data: list = zntrack.deps()
+    data: list[ase.Atoms] = zntrack.deps()
     config: str = zntrack.params_path()
     validation_data: list[ase.Atoms] = zntrack.deps()
     model: t.Optional[ApaxBase] = zntrack.deps(None)