# Trainer code for token-wise prediction model

**URL:** <https://discuss.huggingface.co/t/trainer-code-for-token-wise-prediction-model/18752>\
**Category:** Intermediate\
**Created:** [June 6, 2022, 7:59am UTC](https://discuss.huggingface.co/t/trainer-code-for-token-wise-prediction-model/18752 "2022-06-06T07:59:05Z")\
**Posts on this page:** 1\
**Page:** 1

<div class="post-metadata">

**Author:** ![raghavmallampalli](https://avatars.discourse-cdn.com/v4/letter/r/977dab/32.png) [@raghavmallampalli](https://discuss.huggingface.co/u/raghavmallampalli)\
**Post date:** [June 6, 2022, 7:59am UTC](https://discuss.huggingface.co/t/trainer-code-for-token-wise-prediction-model/18752/1 "2022-06-06T07:59:05Z")

</div>

I wish to train a model that classifies each token in a sequence input. I have a model architecture in place, and I want to use torch.nn.binary\_cross\_entropy\_with\_logits() to calculate loss.

logits = model(inputs)  
inputs shape: (batch\_size, max\_seq\_length)  
logits shape: (batch\_size, max\_seq\_length)

The compute\_loss override example in the docs uses .view(-1) while performing the loss calculation.

This is the code I came up with after reading the docs:

```auto
class CustomTrainer(Trainer):
    def compute_loss(self, model, inputs, return_outputs=False):
        labels = inputs.get("labels")
        outputs = model(**inputs)
        loss = F.binary_cross_entropy_with_logits(
            input=outputs.get("logits").view(-1),
            target=labels.view(-1),
        )
        return (loss, outputs) if return_outputs else loss

```

I have two questions:

1. Can I override compute\_loss in Trainer or should I use Seq2SeqTrainer? The output sequence does not need to be calculated using a beam search.
2. Is .view(-1) still required in the method I intend to use?