Deberta Model Dimension Mismatch
Author: TOP-RX
Pull request: no

### System Info

transformer: 4.24.0
python: 3.8.13

### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

Hello, I am trying to use Deberta_base to do a simple classification, here is my implementation:

Deberta:

```
class DeBertaClassifier(nn.Module):
    
    def __init__(self, num_classes, dropout, pretrain):

        super(DeBertaClassifier, self).__init__()

        self.bert = DebertaModel.from_pretrained(pretrain)
        self.dropout = nn.Dropout(dropout)
        self.linear = nn.Linear(self.bert.config.hidden_size, num_classes)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(dropout)

    def forward(self, input_id, mask):

        deberta_output = self.bert(input_ids=input_id, attention_mask=mask, return_dict=False, output_hidden_states=True)
        hidden_state = deberta_output[0]                    
        pooled_output = hidden_state[:, 0] 
        dropout_output = self.dropout(pooled_output)
        linear_output = self.linear(dropout_output)

        return linear_output
```
Training function:

```
def train(model, train_X, train_y, val_X, val_y, learning_rate, epochs, batch_size, adam_epsilon, model_dir):

    train, val = Dataset(train_X, train_y), Dataset(val_X, val_y)


    train_dataloader = torch.utils.data.DataLoader(train, batch_size=batch_size, shuffle=True)
    val_dataloader = torch.utils.data.DataLoader(val, batch_size=batch_size)

    use_cuda = torch.cuda.is_available()
    device = torch.device("cuda:3" if use_cuda else "cpu")

    criterion = nn.CrossEntropyLoss()
    optimizer = AdamW(model.parameters(), lr=learning_rate, eps=adam_epsilon)

    model = model.to(device)
    criterion = criterion.to(device)

    valid_acc = 0

    for epoch_num in range(epochs):

        total_acc_train = 0
        total_loss_train = 0

        for train_input, train_label in tqdm(train_dataloader):

            train_label = train_label.to(device)
            mask = train_input['attention_mask'].to(device)
            input_id = train_input['input_ids'].squeeze(1).to(device)

            output = model(input_id, mask)

            batch_loss = criterion(output, train_label)
            total_loss_train += batch_loss.item()

            acc = (output.argmax(dim=1) == train_label).sum().item()
            total_acc_train += acc

            model.zero_grad()
            batch_loss.backward()
            optimizer.step()

        total_acc_val = 0
        total_loss_val = 0

        with torch.no_grad():

            for val_input, val_label in val_dataloader:
                val_label = val_label.to(device)
                mask = val_input['attention_mask'].to(device)
                input_id = val_input['input_ids'].squeeze(1).to(device)

                output = model(input_id, mask)

                batch_loss = criterion(output, val_label)
                total_loss_val += batch_loss.item()

                acc = (output.argmax(dim=1) == val_label).sum().item()
                total_acc_val += acc

        print(
            f'Epochs: {epoch_num + 1} | Train Loss: {total_loss_train / len(train_X): .4f} \
                | Train Accuracy: {total_acc_train / len(train_X): .4f} \
                | Val Loss: {total_loss_val / len(val_X): .4f} \
                | Val Accuracy: {total_acc_val / len(val_X): .4f}')
```

```
model_path = r'/Deberta_base'
tokenizer = AutoTokenizer.from_pretrained(model_path) # the downloaded model
X = [tokenizer(text, padding='max_length', max_length=128,
                truncation=True, return_tensors="pt")
        for text in node_text_list]
# Train the model
model = DeBertaClassifier(num_classes=torch.max(y)+1, dropout=args.dropout, pretrain=args.pretrain)

train(model, [X[i] for i in train_idx], y[train_idx], [X[i] for i in valid_idx], y[valid_idx],
      args.learning_rate, args.epochs, args.batch_size, args.adam_epsilon, args.model_dir)

```
I have checked the dimension of input_ids is [128, 1, 128], and attentions is [128, 128] given batch size =128. However, it reports the following bugs:

```
Cell In[28], line 31, in train(model, train_X, train_y, val_X, val_y, learning_rate, epochs, batch_size, adam_epsilon, model_dir)
     28 mask = train_input['attention_mask'].to(device)
     29 input_id = train_input['input_ids'].squeeze(1).to(device)
---> 31 output = model(input_id, mask)
     33 batch_loss = criterion(output, train_label)
     34 total_loss_train += batch_loss.item()

File [~/anaconda3/envs/giant-xrt/lib/python3.8/site-packages/torch/nn/modules/module.py:1130](https://vscode-remote+ssh-002dremote-002bgr.vscode-resource.vscode-cdn.net/egr/research-dselab/hanhaoy1/proc_data_xrt/LM/~/anaconda3/envs/giant-xrt/lib/python3.8/site-packages/torch/nn/modules/module.py:1130), in Module._call_impl(self, *input, **kwargs)
   1126 # If we don't have any hooks, we want to skip the rest of the logic in
   1127 # this function, and just call forward.
...
--> 826     embeddings = embeddings * mask
    828 embeddings = self.dropout(embeddings)
    829 return embeddings

RuntimeError: The size of tensor a (768) must match the size of tensor b (128) at non-singleton dimension 2
```
The bug happened at `deberta_output = self.bert(input_ids=input_id, attention_mask=mask, return_dict=False, output_hidden_states=True)`
I am wondering if I missed something. It won't happen when I use Bert. Any help would be appreciated.

### Expected behavior

Dimensions match correctly.

Comments: