diff --git a/mala/descriptors/atomic_density.py b/mala/descriptors/atomic_density.py
index 558f6695..45c889c3 100755
--- a/mala/descriptors/atomic_density.py
+++ b/mala/descriptors/atomic_density.py
@@ -213,6 +213,7 @@ def __calculate_lammps(self, outdir, **kwargs):
             array_shape=(nrows_ggrid, ncols_ggrid),
             use_fp64=use_fp64,
         )
+
         self._clean_calculation(lmp, keep_logs)
 
         # In comparison to bispectrum, the atomic density always returns
@@ -241,7 +242,7 @@ def __calculate_lammps(self, outdir, **kwargs):
                         self.grid_dimensions[2],
                         self.grid_dimensions[1],
                         self.grid_dimensions[0],
-                        7,
+                        6 + len(set(self._atoms.numbers)),
                     )
                 )
                 gaussian_descriptors_np = gaussian_descriptors_np.transpose(
diff --git a/mala/targets/density.py b/mala/targets/density.py
index 61623fe2..be04f6f5 100644
--- a/mala/targets/density.py
+++ b/mala/targets/density.py
@@ -1214,7 +1214,7 @@ def __setup_total_energy_module(
             t0 = time.perf_counter()
             gaussian_descriptors = np.reshape(
                 gaussian_descriptors,
-                [number_of_gridpoints_mala, 1],
+                [number_of_gridpoints_mala, len(set(self.atoms.numbers))],
                 order="F",
             )
             reference_gaussian_descriptors = np.reshape(
@@ -1244,7 +1244,7 @@ def __setup_total_energy_module(
                 reference_gaussian_descriptors,
                 sigma,
                 number_of_gridpoints,
-                1,
+                len(set(self.atoms.numbers)),
             )
             barrier()
             printout(