return_type

BiomedSciAI · Feb 11, 2024 · e423e0a · e423e0a
1 parent 8f25995
commit e423e0a
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/fuse/data/utils/collates.py b/fuse/data/utils/collates.py
@@ -192,7 +192,9 @@ def _pad_size(value: torch.Tensor, dim: int) -> List[int]:
  return default_collate(padded_values)
 
  @staticmethod
- def crop_padding(input_ids_list: List[torch.Tensor], pad_token_id: int):
+ def crop_padding(
+ input_ids_list: List[torch.Tensor], pad_token_id: int
+ ) -> torch.Tensor:
  """
  Crop padding of a batch of input_ids tensors to the minimum length possible.