diff --git a/EMRI_DET/validate.py b/EMRI_DET/validate.py
index 190a63f8576bed3cb2ff656642f0077e43ea339e..a6476ef05cfee25b8f37d78cfb36c2ccd64db5f4 100644
--- a/EMRI_DET/validate.py
+++ b/EMRI_DET/validate.py
@@ -216,7 +216,7 @@ def grid_heatmap_corner(dataframe, truth_column, pred_column, log=True, ratio=Fa
     vmin = np.min([np.min(hmap) for hmap in plotmaps])
     vmax = np.max([np.max(hmap) for hmap in plotmaps])
 
-    fig, ax = plt.subplots(ncols=5, nrows=5, figsize=(8 + int(2*nparams),6 + int(2*nparams)))
+    fig, ax = plt.subplots(ncols=nparams, nrows=nparams, figsize=(8 + int(2*nparams),6 + int(2*nparams)))
     plt.subplots_adjust(wspace=0.05, hspace=0.05)
     num = 0
     for i in range(nparams):