Mismatch of tensor shapes in CrossEntropyLoss for custom head layer in BART
Author: CodingSaturn
Pull request: no

Hi,

so far I've been working with the BartForConditionalGeneration. Now I want to use a custom head layer instead. 
In a linear layer after the base models decoder, I want to input the output of the base bart model and additional some numerical data similar to the code [here](https://colab.research.google.com/drive/1ZLfcB16Et9U2V-udrw8zwrfChFCIhomz?usp=sharing#scrollTo=m-TTyOMJOGBD). Following this I came up with the following forward function:

```
def forward(self, input_ids, tokens, **kwargs):
    labels = kwargs.get('labels')
    attn_mask = kwargs.get('attention_mask')
    out = self.model_base(input_ids, attention_mask=attn_mask)
    token_features = tokens.unsqueeze(1)
    concat= torch.concat((out[0][:, 0, :], token_features), dim=-1)
    out_lin = self.custom_layer(concat)

    loss_fct = torch.nn.CrossEntropyLoss()
    masked_lm_loss = loss_fct(out_lin.view(-1, self.model.config.vocab_size), labels.view(-1))
```

where out_lin is the following linear layer:
```
self.custom_layer = torch.nn.Linear(in_features = self.hidden_dim + self.token_dim, out_features = self.model.config.vocab_size)
```

For the loss function I took orientation from the original code for the [BartForConditionalGeneration](https://huggingface.co/transformers/v2.11.0/_modules/transformers/modeling_bart.html#BartForConditionalGeneration):
```
outputs = self.model(
            input_ids,
            attention_mask=attention_mask,
            decoder_input_ids=decoder_input_ids,
            encoder_outputs=encoder_outputs,
            decoder_attention_mask=decoder_attention_mask,
            decoder_cached_states=decoder_cached_states,
            use_cache=use_cache,
        )
        lm_logits = F.linear(outputs[0], self.model.shared.weight, bias=self.final_logits_bias)
        outputs = (lm_logits,) + outputs[1:]  # Add cache, hidden states and attention if they are here
        if lm_labels is not None:
            loss_fct = nn.CrossEntropyLoss()
            # TODO(SS): do we need to ignore pad tokens in lm_labels?
            masked_lm_loss = loss_fct(lm_logits.view(-1, self.config.vocab_size), lm_labels.view(-1))
            outputs = (masked_lm_loss,) + outputs

        return outputs
```

The error I obtain is
```
---------------------------------------------------------------------------
File c:\Users\M\Anaconda\envs\simp_env\lib\site-packages\pytorch_lightning\trainer\trainer.py:579, in Trainer.fit(self, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path)
    577     raise TypeError(f"`Trainer.fit()` requires a `LightningModule`, got: {model.__class__.__qualname__}")
    578 self.strategy._lightning_module = model
--> 579 call._call_and_handle_interrupt(
    580     self, self._fit_impl, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path
    581 )

File c:\Users\M\Anaconda\envs\simp_env\lib\site-packages\pytorch_lightning\trainer\call.py:38, in _call_and_handle_interrupt(trainer, trainer_fn, *args, **kwargs)
     36         return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
     37     else:
---> 38         return trainer_fn(*args, **kwargs)
     40 except _TunerExitException:
...
   3024 if size_average is not None or reduce is not None:
   3025     reduction = _Reduction.legacy_get_string(size_average, reduce)
-> 3026 return torch._C._nn.cross_entropy_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index, label_smoothing)

ValueError: Expected input batch_size (6) to match target batch_size (1536).
```

So far I understand that 6 is the batch size of my data output from the custom linear layer (which is torch.Size([6, 50267]) where 50267 is the self.final_logits_bias/vocab_size). My labels have the shape torch.Size([6, 256]) which when flattened leads to the 1536.
As my labels have the same shape as before and my layer seems to me the same as the one from the ConditionalGenerationModel which I used before, I am unsure why I suddenly receive this size incompatibility issue, when I did not before.

Furthermore, I am unsure why the first code referenced here uses only hidden_states.last_hidden_state[:, 0, :] so only batch_size and hidden_size but not the sequence length. Without it my data has the shape torch.Size([6, 256, 768]).

I would be thankful for any guidance on how to make the tensors compatible and make the custom layer work.
Did I misunderstand the examples mentioned above?

Comments:
You should try the [forums](https://discuss.huggingface.co/) for questions like this as we keep issues for bugs and feature requests only.

Oh I am very sorry, I moved the question to the forum. Thank your for the hint!