# Buggy code
def forward(self, model, sample, reduce=True):
    """Compute the loss for the given sample.

    Returns a tuple with three elements:
    1) the loss
    2) the sample size, which is used as the denominator for the gradient
    3) logging outputs to display while training
    """
    assert (
        hasattr(model, "classification_heads")
        and self.args.classification_head_name in model.classification_heads
    ), (
        "model must provide sentence classification head for --criterion=sentence_prediction"
    )

    logits, _ = model(
        **sample["net_input"],
        features_only=True,
        classification_head_name=self.args.classification_head_name,
    )
    targets = model.get_targets(sample, [logits]).view(-1)
    sample_size = targets.numel()

    if not self.args.regression_target:
        loss = F.nll_loss(
            F.log_softmax(logits, dim=-1, dtype=torch.float32),
            targets,
            reduction="sum",
        )
    else:
        logits = logits.squeeze().float()
        targets = targets.float()
        loss = F.mse_loss(
            logits,
            targets,
            reduction="sum",
        )

    logging_output = {
        "loss": loss.data,
        "ntokens": sample["ntokens"],
        "nsentences": sample_size,
        "sample_size": sample_size,
    }
    if not self.args.regression_target:
        preds = logits.argmax(dim=1)
        logging_output["ncorrect"] = utils.item((preds == targets).sum())

    return loss, sample_size, logging_output


# Traceback (AttributeError)
2020-03-08 13:40:11 | INFO | fairseq_cli.train | model roberta_large, criterion SentencePredictionCriterion
2020-03-08 13:40:11 | INFO | fairseq_cli.train | num. model params: 357499983 (num. trained: 357499983)
2020-03-08 13:40:11 | INFO | fairseq_cli.train | training on 1 GPUs
2020-03-08 13:40:11 | INFO | fairseq_cli.train | max tokens per GPU = 4 and max sentences per GPU = 2
2020-03-08 13:40:11 | INFO | fairseq.models.roberta.model | Overwriting classification_heads.head_name.dense.weight
2020-03-08 13:40:11 | INFO | fairseq.models.roberta.model | Overwriting classification_heads.head_name.dense.bias
2020-03-08 13:40:11 | INFO | fairseq.models.roberta.model | Overwriting classification_heads.head_name.out_proj.weight
2020-03-08 13:40:11 | INFO | fairseq.models.roberta.model | Overwriting classification_heads.head_name.out_proj.bias
2020-03-08 13:40:11 | INFO | fairseq.trainer | loaded checkpoint roberta.large/model.pt (epoch 1 @ 0 updates)
2020-03-08 13:40:11 | INFO | fairseq.trainer | loading train data for epoch 1
2020-03-08 13:40:11 | INFO | fairseq.data.data_utils | loaded 160600 examples from: ../data/input0/train
2020-03-08 13:40:11 | INFO | fairseq.data.data_utils | loaded 160600 examples from: ../data/label/train
2020-03-08 13:40:11 | INFO | fairseq.tasks.sentence_prediction | Loaded train with #samples: 160600
2020-03-08 13:40:11 | WARNING | fairseq.data.data_utils | 160580 samples have invalid sizes and will be skipped, max_positions=4, first few sample ids=[141335, 12811, 122058, 104432, 1925, 141500, 65517, 143950, 59283, 155828]
Traceback (most recent call last):
File "train.py", line 11, in <module>
cli_main()
File "/nas/home/thawani/MCS/fairseq/fairseq_cli/train.py", line 322, in cli_main
main(args)
File "/nas/home/thawani/MCS/fairseq/fairseq_cli/train.py", line 100, in main
train(args, trainer, task, epoch_itr)
File "/nas/home/thawani/anaconda3/envs/env/lib/python3.7/contextlib.py", line 74, in inner
return func(*args, **kwds)
File "/nas/home/thawani/MCS/fairseq/fairseq_cli/train.py", line 177, in train
log_output = trainer.train_step(samples)
File "/nas/home/thawani/anaconda3/envs/env/lib/python3.7/contextlib.py", line 74, in inner
return func(*args, **kwds)
File "/nas/home/thawani/MCS/fairseq/fairseq/trainer.py", line 319, in train_step
ignore_grad=is_dummy_batch,
File "/nas/home/thawani/MCS/fairseq/fairseq/tasks/fairseq_task.py", line 337, in train_step
loss, sample_size, logging_output = criterion(model, sample)
File "/nas/home/thawani/anaconda3/envs/env/lib/python3.7/site-packages/torch/nn/modules/module.py", line 532, in __call__
result = self.forward(*input, **kwargs)
File "/nas/home/thawani/MCS/fairseq/fairseq/criterions/sentence_prediction.py", line 40, in forward
and self.args.classification_head_name in model.classification_heads
File "/nas/home/thawani/anaconda3/envs/env/lib/python3.7/site-packages/torch/nn/modules/module.py", line 576, in __getattr__
type(self).__name__, name))
AttributeError: 'SentencePredictionCriterion' object has no attribute 'args'

# Fixed code
def forward(self, model, sample, reduce=True):
    """Compute the loss for the given sample.

    Returns a tuple with three elements:
    1) the loss
    2) the sample size, which is used as the denominator for the gradient
    3) logging outputs to display while training
    """
    assert (
        hasattr(model, "classification_heads")
        and self.classification_head_name in model.classification_heads
    ), (
        "model must provide sentence classification head for --criterion=sentence_prediction"
    )

    logits, _ = model(
        **sample["net_input"],
        features_only=True,
        classification_head_name=self.classification_head_name,
    )
    targets = model.get_targets(sample, [logits]).view(-1)
    sample_size = targets.numel()

    if not self.regression_target:
        loss = F.nll_loss(
            F.log_softmax(logits, dim=-1, dtype=torch.float32),
            targets,
            reduction="sum",
        )
    else:
        logits = logits.squeeze().float()
        targets = targets.float()
        loss = F.mse_loss(
            logits,
            targets,
            reduction="sum",
        )

    logging_output = {
        "loss": loss.data,
        "ntokens": sample["ntokens"],
        "nsentences": sample_size,
        "sample_size": sample_size,
    }
    if not self.regression_target:
        preds = logits.argmax(dim=1)
        logging_output["ncorrect"] = utils.item((preds == targets).sum())

    return loss, sample_size, logging_output