diff --git a/dreamer4/dreamer4.py b/dreamer4/dreamer4.py
index 71dd3c6..f44793e 100644
--- a/dreamer4/dreamer4.py
+++ b/dreamer4/dreamer4.py
@@ -1298,6 +1298,7 @@ class DynamicsModel(Module):
                     latents = noised_latent_with_context,
                     signal_levels = signal_levels_with_context,
                     step_sizes = step_size,
+                    rewards = decoded_rewards,
                     latent_is_noised = True,
                     return_pred_only = True,
                     return_agent_tokens = True