Fix type casting issue in mask length calculation

pytorch · Sep 7, 2023 · 0c0f3ea · 0c0f3ea
1 parent ede4309
commit 0c0f3ea
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torchaudio/models/wav2vec2/components.py b/torchaudio/models/wav2vec2/components.py
@@ -894,7 +894,7 @@ def _compute_mask_indices(
         if mask_type == "static":
             lengths = torch.full((num_mask,), mask_length)
         elif mask_type == "uniform":
-            lengths = torch.randint(mask_other, mask_length * 2 + 1, size=(num_mask,))
+            lengths = torch.randint(int(mask_other), mask_length * 2 + 1, size=(num_mask,))
         elif mask_type == "normal":
             lengths = torch.normal(mask_length, mask_other, size=(num_mask,))
             lengths = torch.maximum(torch.ones(1), torch.round(lengths)).int()