diff --git a/src/lava/proc/dense/models.py b/src/lava/proc/dense/models.py
index 5bc69ce69..f9b833976 100644
--- a/src/lava/proc/dense/models.py
+++ b/src/lava/proc/dense/models.py
@@ -212,11 +212,13 @@ def calc_act(self, s_in) -> np.ndarray:
         #  which is then transposed to get the activation matrix.
         return np.reshape(
             np.sum(self.get_delay_wgts_mat(self.weights,
-                                           self.delays) * s_in, axis=1),
-            (np.max(self.delays) + 1, self.weights.shape[0])).T
+                                           self.delays,
+                                           self.a_buff.shape[-1] - 1) * s_in,
+                   axis=1),
+            (self.a_buff.shape[-1], self.weights.shape[0])).T
 
     @staticmethod
-    def get_delay_wgts_mat(weights, delays) -> np.ndarray:
+    def get_delay_wgts_mat(weights, delays, max_delay) -> np.ndarray:
         """
         Create a matrix where the synaptic weights are separated
         by their corresponding delays. The first matrix contains all the
@@ -235,7 +237,7 @@ def get_delay_wgts_mat(weights, delays) -> np.ndarray:
         """
         return np.vstack([
             np.where(delays == k, weights, 0)
-            for k in range(np.max(delays) + 1)
+            for k in range(max_delay + 1)
         ])
 
     def update_act(self, s_in):