# Buggy code
def forward(
    self,
    token_ids: torch.LongTensor,
    mask: torch.BoolTensor,
    offsets: torch.LongTensor,
    wordpiece_mask: torch.BoolTensor,
    type_ids: Optional[torch.LongTensor] = None,
    segment_concat_mask: Optional[torch.BoolTensor] = None,
) -> torch.Tensor:  # type: ignore
    """
    # Parameters

    token_ids: `torch.LongTensor`
        Shape: [batch_size, num_wordpieces] (for exception see `PretrainedTransformerEmbedder`).
    mask: `torch.BoolTensor`
        Shape: [batch_size, num_orig_tokens].
    offsets: `torch.LongTensor`
        Shape: [batch_size, num_orig_tokens, 2].
        Maps indices for the original tokens, i.e. those given as input to the indexer,
        to a span in token_ids. `token_ids[i][offsets[i][j][0]:offsets[i][j][1] + 1]`
        corresponds to the original j-th token from the i-th batch.
    wordpiece_mask: `torch.BoolTensor`
        Shape: [batch_size, num_wordpieces].
    type_ids: `Optional[torch.LongTensor]`
        Shape: [batch_size, num_wordpieces].
    segment_concat_mask: `Optional[torch.BoolTensor]`
        See `PretrainedTransformerEmbedder`.

    # Returns

    `torch.Tensor`
        Shape: [batch_size, num_orig_tokens, embedding_size].
    """
    # Shape: [batch_size, num_wordpieces, embedding_size].
    embeddings = self._matched_embedder(
        token_ids,
        wordpiece_mask,
        type_ids=type_ids,
        segment_concat_mask=segment_concat_mask,
    )

    # span_embeddings: (batch_size, num_orig_tokens, max_span_length, embedding_size)
    # span_mask: (batch_size, num_orig_tokens, max_span_length)
    span_embeddings, span_mask = util.batched_span_select(
        embeddings.contiguous(), offsets
    )
    span_mask = span_mask.unsqueeze(-1)
    span_embeddings *= span_mask  # zero out paddings

    span_embeddings_sum = span_embeddings.sum(2)
    span_embeddings_len = span_mask.sum(2)
    # Shape: (batch_size, num_orig_tokens, embedding_size)
    orig_embeddings = span_embeddings_sum / span_embeddings_len

    return orig_embeddings


# Traceback (TypeError)
Traceback (most recent call last):
File "/home/arthur/question-generation/models/sg_dqg.py", line 156, in <module>
preprocess(args.ds)
File "/home/arthur/question-generation/models/sg_dqg.py", line 124, in preprocess
coreferences = coreference_resolution(evidences_list)
File "/home/arthur/question-generation/models/sg_dqg.py", line 74, in coreference_resolution
document="Besides its prominence in sports, Notre Dame is also a large, four-year, highly residential research University, and is consistently ranked among the top twenty universities in the United States  and as a major global university."
File "/home/arthur/anaconda3/envs/tensorflow_env/lib/python3.6/site-packages/allennlp_models/coref/coref_predictor.py", line 65, in predict
return self.predict_json({"document": document})
File "/home/arthur/anaconda3/envs/tensorflow_env/lib/python3.6/site-packages/allennlp/predictors/predictor.py", line 48, in predict_json
return self.predict_instance(instance)
File "/home/arthur/anaconda3/envs/tensorflow_env/lib/python3.6/site-packages/allennlp/predictors/predictor.py", line 171, in predict_instance
outputs = self._model.forward_on_instance(instance)
File "/home/arthur/anaconda3/envs/tensorflow_env/lib/python3.6/site-packages/allennlp/models/model.py", line 142, in forward_on_instance
return self.forward_on_instances([instance])[0]
File "/home/arthur/anaconda3/envs/tensorflow_env/lib/python3.6/site-packages/allennlp/models/model.py", line 167, in forward_on_instances
model_input = util.move_to_device(dataset.as_tensor_dict(), cuda_device)
File "/home/arthur/anaconda3/envs/tensorflow_env/lib/python3.6/site-packages/allennlp/data/batch.py", line 139, in as_tensor_dict
for field, tensors in instance.as_tensor_dict(lengths_to_use).items():
File "/home/arthur/anaconda3/envs/tensorflow_env/lib/python3.6/site-packages/allennlp/data/instance.py", line 99, in as_tensor_dict
tensors[field_name] = field.as_tensor(padding_lengths[field_name])
File "/home/arthur/anaconda3/envs/tensorflow_env/lib/python3.6/site-packages/allennlp/data/fields/text_field.py", line 103, in as_tensor
self._indexed_tokens[indexer_name], indexer_lengths[indexer_name]
File "/home/arthur/anaconda3/envs/tensorflow_env/lib/python3.6/site-packages/allennlp/data/token_indexers/pretrained_transformer_mismatched_indexer.py", line 96, in as_padded_tensor_dict
offsets_tokens, offsets_padding_lengths, default_value=lambda: (0, 0)
TypeError: not a sequence

# Fixed code
def forward(
    self,
    token_ids: torch.LongTensor,
    mask: torch.BoolTensor,
    offsets: torch.LongTensor,
    wordpiece_mask: torch.BoolTensor,
    type_ids: Optional[torch.LongTensor] = None,
    segment_concat_mask: Optional[torch.BoolTensor] = None,
) -> torch.Tensor:  # type: ignore
    """
    # Parameters

    token_ids: `torch.LongTensor`
        Shape: [batch_size, num_wordpieces] (for exception see `PretrainedTransformerEmbedder`).
    mask: `torch.BoolTensor`
        Shape: [batch_size, num_orig_tokens].
    offsets: `torch.LongTensor`
        Shape: [batch_size, num_orig_tokens, 2].
        Maps indices for the original tokens, i.e. those given as input to the indexer,
        to a span in token_ids. `token_ids[i][offsets[i][j][0]:offsets[i][j][1] + 1]`
        corresponds to the original j-th token from the i-th batch.
    wordpiece_mask: `torch.BoolTensor`
        Shape: [batch_size, num_wordpieces].
    type_ids: `Optional[torch.LongTensor]`
        Shape: [batch_size, num_wordpieces].
    segment_concat_mask: `Optional[torch.BoolTensor]`
        See `PretrainedTransformerEmbedder`.

    # Returns

    `torch.Tensor`
        Shape: [batch_size, num_orig_tokens, embedding_size].
    """
    # Shape: [batch_size, num_wordpieces, embedding_size].
    embeddings = self._matched_embedder(
        token_ids,
        wordpiece_mask,
        type_ids=type_ids,
        segment_concat_mask=segment_concat_mask,
    )

    # span_embeddings: (batch_size, num_orig_tokens, max_span_length, embedding_size)
    # span_mask: (batch_size, num_orig_tokens, max_span_length)
    span_embeddings, span_mask = util.batched_span_select(
        embeddings.contiguous(), offsets
    )
    span_mask = span_mask.unsqueeze(-1)
    span_embeddings *= span_mask  # zero out paddings

    span_embeddings_sum = span_embeddings.sum(2)
    span_embeddings_len = span_mask.sum(2)
    # Shape: (batch_size, num_orig_tokens, embedding_size)
    orig_embeddings = span_embeddings_sum / span_embeddings_len

    # All the places where the span length is zero, write in zeros.
    orig_embeddings[(span_embeddings_len == 0).expand(orig_embeddings.shape)] = 0

    return orig_embeddings