code: def lengths_to_mask(*lengths, **kwargs):
      
      # Squeeze to deal with random additional dimensions
      lengths = [l.squeeze().tolist() if torch.is_tensor(l) else l for l in lengths]
  
      # For cases where length is a scalar, this needs to convert it to a list.
      lengths = [l if isinstance(l, list) else [l] for l in lengths]
      assert all(len(l) == len(lengths[0]) for l in lengths)
      batch_size = len(lengths[0])
      other_dimensions = tuple([int(max(l)) for l in lengths])
      mask = torch.zeros(batch_size, *other_dimensions, **kwargs)
      for i, length in enumerate(zip(*tuple(lengths))):
          mask[i][[slice(int(l)) for l in length]].fill_(1)
      return mask.byte()
docs: Given a list of lengths, create a batch mask.
  
      Example:
          >>> lengths_to_mask([1, 2, 3])
          tensor([[1, 0, 0],
                  [1, 1, 0],
                  [1, 1, 1]], dtype=torch.uint8)
          >>> lengths_to_mask([1, 2, 2], [1, 2, 2])
          tensor([[[1, 0],
                   [0, 0]],
          <BLANKLINE>
                  [[1, 1],
                   [1, 1]],
          <BLANKLINE>
                  [[1, 1],
                   [1, 1]]], dtype=torch.uint8)
  
      Args:
          *lengths (list of int or torch.Tensor)
          **kwargs: Keyword arguments passed to ``torch.zeros`` upon initially creating the returned
            tensor.
  
      Returns:
          torch.ByteTensor