From b5744237bf144b062bb13774834a6a387d7a9ac0 Mon Sep 17 00:00:00 2001
From: lucidrains <lucidrains@gmail.com>
Date: Wed, 8 Oct 2025 06:58:46 -0700
Subject: [PATCH] fix

---
 dreamer4/dreamer4.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/dreamer4/dreamer4.py b/dreamer4/dreamer4.py
index b5c8cd0..1b2a274 100644
--- a/dreamer4/dreamer4.py
+++ b/dreamer4/dreamer4.py
@@ -1175,6 +1175,7 @@ class DynamicsModel(Module):
         # each agent token will have the reward embedding of the previous time step - but could eventually just give reward its own token
 
         self.add_reward_embed_to_agent_token = add_reward_embed_to_agent_token
+        self.add_reward_embed_dropout = add_reward_embed_dropout
 
         self.reward_encoder = SymExpTwoHot(
             **reward_encoder_kwargs,