From b62c08be65d70c1b49a5276ed10054339c5a1884 Mon Sep 17 00:00:00 2001
From: lucidrains <lucidrains@gmail.com>
Date: Wed, 8 Oct 2025 08:42:25 -0700
Subject: [PATCH] fix task embed in presence of multiple agent tokens

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

diff --git a/dreamer4/dreamer4.py b/dreamer4/dreamer4.py
index 18e22f8..091be87 100644
--- a/dreamer4/dreamer4.py
+++ b/dreamer4/dreamer4.py
@@ -1519,7 +1519,7 @@ class DynamicsWorldModel(Module):
             assert self.num_tasks > 0
 
             task_embeds = self.task_embed(tasks)
-            agent_tokens = agent_tokens + task_embeds
+            agent_tokens = einx.add('b ... d, b d', agent_tokens, task_embeds)
 
         # maybe evolution