Return BoolTensor from get_mask_from_sequence_lengths
@@ -100,7 +100,7 @@ def batch_tensor_dicts(
return batched_tensors
-def get_lengths_from_binary_sequence_mask(mask: torch.BoolTensor):
+def get_lengths_from_binary_sequence_mask(mask: torch.BoolTensor) -> torch.LongTensor:
"""
Compute sequence lengths for each batch element in a tensor using a
binary mask.
@@ -119,7 +119,9 @@ def get_lengths_from_binary_sequence_mask(mask: torch.BoolTensor):
return mask.sum(-1)
-def get_mask_from_sequence_lengths(sequence_lengths: torch.Tensor, max_length: int) -> torch.Tensor:
+def get_mask_from_sequence_lengths(
+ sequence_lengths: torch.Tensor, max_length: int
+) -> torch.BoolTensor:
"""
Given a variable of shape `(batch_size,)` that represents the sequence lengths of each batch
element, this function returns a `(batch_size, max_length)` mask variable. For example, if
@@ -133,7 +135,7 @@ def get_mask_from_sequence_lengths(sequence_lengths: torch.Tensor, max_length: i
# (batch_size, max_length)
ones = sequence_lengths.new_ones(sequence_lengths.size(0), max_length)
range_tensor = ones.cumsum(dim=1)
- return (sequence_lengths.unsqueeze(1) >= range_tensor).long()
+ return sequence_lengths.unsqueeze(1) >= range_tensor
def sort_batch_by_length(tensor: torch.Tensor, sequence_lengths: torch.Tensor):
