diff --git a/setup.py b/setup.py
index 04dd15b..9562d6c 100644
--- a/setup.py
+++ b/setup.py
@@ -2,7 +2,7 @@
 
 setup(
     name="tensorpotential",
-    version="0.4.2",
+    version="0.4.3",
     packages=find_packages(include=["tensorpotential", "tensorpotential.*"]),
     url="https://github.com/ICAMS/grace-tensorpotential",
     license="Academic Software License (ASL)",
diff --git a/tensorpotential/data/databuilder.py b/tensorpotential/data/databuilder.py
index 14157ea..ae80146 100644
--- a/tensorpotential/data/databuilder.py
+++ b/tensorpotential/data/databuilder.py
@@ -413,7 +413,9 @@ def pad_batch(self, batch, max_pad_dict):
         pad_nat, pad_nneigh, pad_nstruct = get_padding_dims(batch, max_pad_dict)
 
         max_nat = np.array(max_pad_dict[constants.PAD_MAX_N_ATOMS]).astype(np.int32)
-        max_structs = max_pad_dict[constants.PAD_MAX_N_STRUCTURES].astype(np.int32)
+        max_structs = np.array(max_pad_dict[constants.PAD_MAX_N_STRUCTURES]).astype(
+            np.int32
+        )
 
         batch[constants.N_ATOMS_BATCH_TOTAL] = max_nat
         batch[constants.N_STRUCTURES_BATCH_TOTAL] = max_structs