From 77724049e2655378778e34e5a67a7c94ef561c67 Mon Sep 17 00:00:00 2001 From: lucidrains Date: Mon, 6 Oct 2025 09:43:16 -0700 Subject: [PATCH] fix latent / modality attention pattern in video tokenizer, thanks to another researcher --- dreamer4/dreamer4.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/dreamer4/dreamer4.py b/dreamer4/dreamer4.py index f21027b..fff3b4f 100644 --- a/dreamer4/dreamer4.py +++ b/dreamer4/dreamer4.py @@ -896,7 +896,7 @@ class VideoTokenizer(Module): # modality can only attend to itself while latents can attend to everything # similar to agent token in dynamics model - encoder_attend_fn = get_attend_fn(use_flex, seq_len, seq_len, special_attend_only_itself = True) + encoder_attend_fn = get_attend_fn(use_flex, seq_len, seq_len, special_attend_only_itself = False) # encoder @@ -937,7 +937,7 @@ class VideoTokenizer(Module): # decoder attend - decoder_attend_fn = get_attend_fn(use_flex, seq_len, seq_len) + decoder_attend_fn = get_attend_fn(use_flex, seq_len, seq_len, special_attend_only_itself = True) # decoder attention