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

cleanup

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

4
dreamer4/dreamer4.py
View File

@ -450,9 +450,7 @@ class VideoTokenizer(Module):
if mask_patches:
min_mask_prob, max_mask_prob = self.per_image_patch_mask_prob
uniform_prob = torch.rand(tokens.shape[:2], device = tokens.device) # (b t)
mask_prob = uniform_prob * (max_mask_prob - min_mask_prob) + min_mask_prob
mask_prob = torch.empty(tokens.shape[:2], device = tokens.device).uniform_(min_mask_prob, max_mask_prob) # (b t)
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.