hongshaorou/dreamer4
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

ellipsis

Browse Source
This commit is contained in:
lucidrains 2025-10-02 09:14:39 -07:00
parent e6c808960f
commit 0b503d880d
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
dreamer4/dreamer4.py
View File

@ -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