fix type annotation...

graphcore-research · Sep 18, 2023 · 982ca57 · 982ca57
1 parent fc8ac74
commit 982ca57
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/pyscf_ipu/nanoDFT/nanoDFT.py b/pyscf_ipu/nanoDFT/nanoDFT.py
@@ -191,8 +191,8 @@ def _nanoDFT(state, opts, mol):
 FloatNxN = Float[Array, "N N"]
 FloatNxNxNxN = Float[Array, "N N N N"]
 Grid = Float[Array, "4 grid_size N"]
-FloatArray = Float[Array]
-IntArray = Int[Array]
+FloatArray = Float[Array, "..."]
+IntArray = Int[Array, "..."]
 
 
 @chex.dataclass