diff --git a/adv_control/utils.py b/adv_control/utils.py
index edd45a8..be6b73b 100644
--- a/adv_control/utils.py
+++ b/adv_control/utils.py
@@ -523,7 +523,10 @@ class AdvancedControlBase:
                         if o[i] is None:
                             o[i] = prev_val
                         else:
-                            o[i] += prev_val
+                            if o[i].shape[0] < prev_val.shape[0]:
+                                o[i] = prev_val + o[i]
+                            else:
+                                o[i] += prev_val
         return out
 
     def prepare_mask_cond_hint(self, x_noisy: Tensor, t, cond, batched_number, dtype=None):