roar

lucidrains · Aug 23, 2023 · 35bf3c0 · 35bf3c0
1 parent fa1b7d2
commit 35bf3c0
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/soundstorm_pytorch/soundstorm.py b/soundstorm_pytorch/soundstorm.py
@@ -79,6 +79,7 @@ def coin_flip():
 
 # tensor helpers
 
+@beartype
 def get_mask_subset_prob(
     mask: Tensor,
     prob: Union[float, Tensor],