diff --git a/python/snewpy/models/base.py b/python/snewpy/models/base.py
index 6b537250..9a479bcd 100644
--- a/python/snewpy/models/base.py
+++ b/python/snewpy/models/base.py
@@ -303,7 +303,7 @@ def get_initial_spectra(self, t, E, flavors=Flavor):
                     - loggamma(1+a) + a*np.log(E) - (1+a)*(E/Ea)) / (u.erg * u.s)
             #remove bad values
             result[np.isnan(result)] = 0
-            result[E==0] = 0
+            result[:, E==0] = 0
             #remove unnecessary dimensions, if E or t was scalar:
             result = np.squeeze(result)
             initialspectra[flavor] = result