From 0b503d880d2e0b0d3b1ff73c968cd7f96ac331cd Mon Sep 17 00:00:00 2001
From: lucidrains <lucidrains@gmail.com>
Date: Thu, 2 Oct 2025 09:14:39 -0700
Subject: [PATCH] ellipsis

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

diff --git a/dreamer4/dreamer4.py b/dreamer4/dreamer4.py
index 3998a6b..4b274f0 100644
--- a/dreamer4/dreamer4.py
+++ b/dreamer4/dreamer4.py
@@ -457,7 +457,7 @@ class VideoTokenizer(Module):
             mask_prob = repeat(mask_prob, 'b t -> b t vh vw', vh = tokens.shape[2], vw = tokens.shape[3])
             mask_patch = torch.bernoulli(mask_prob) == 1.
 
-            tokens = einx.where('b t vh vw, d, b t vh vw d', mask_patch, self.mask_token, tokens)
+            tokens = einx.where('..., d, ... d', mask_patch, self.mask_token, tokens)
 
         # pack space