title: ValueError: Target size (torch.Size([16, 2])) must be the same as input size (torch.Size([16, 128, 2]))
body: <p>I am doing a text classification using pytorch and RoBERTa. The model will classify the text into 2 classes but I got this error while doing the training.</p>
  <p>ValueError: Target size (torch.Size([16, 2])) must be the same as input size (torch.Size([16, 128, 2])),And max_len=128,batch_size=16.and this is pretrained model <a href="https://huggingface.co/nlp-waseda/roberta-base-japanese" rel="nofollow noreferrer">https://huggingface.co/nlp-waseda/roberta-base-japanese</a></p>
  <pre><code>def train_model(dataset_train, dataset_valid, batch_size, model, criterion, optimizer, num_epochs, device=None):
    &quot;&quot;&quot;モデルの学習を実行し、損失・正解率のログを返す&quot;&quot;&quot;
    # デバイスの指定
    model.to(device)
  
    # dataloaderの作成
    dataloader_train = DataLoader(dataset_train, batch_size=batch_size, shuffle=True)
    dataloader_valid = DataLoader(dataset_valid, batch_size=len(dataset_valid), shuffle=False)
  
    # 学習
    log_train = []
    log_valid = []
    for epoch in range(num_epochs):
      # 開始時刻の記録
      s_time = time.time()
  
      # 訓練モードに設定
      model.train()
      for data in dataloader_train:
        # デバイスの指定
        ids = data['ids'].to(device)
        mask = data['mask'].to(device)
        labels = data['labels'].to(device)
  
        # 勾配をゼロで初期化
        optimizer.zero_grad()
  
        # 順伝播 + 誤差逆伝播 + 重み更新
        outputs = model(ids, mask)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
  
      # 損失と正解率の算出
      loss_train, acc_train = calculate_loss_and_accuracy(model, dataloader_train, device, criterion=criterion)
      loss_valid, acc_valid = calculate_loss_and_accuracy(model, dataloader_valid, device, criterion=criterion)
      log_train.append([loss_train, acc_train])
      log_valid.append([loss_valid, acc_valid])
  
  
      # 終了時刻の記録
      e_time = time.time()
  
      # ログを出力
      print(f'epoch: {epoch + 1}, loss_train: {loss_train:.4f}, accuracy_train: {acc_train:.4f}, loss_valid: {loss_valid:.4f}, accuracy_valid: {acc_valid:.4f}, {(e_time - s_time):.4f}sec') 
  
    return {'train': log_train, 'valid': log_valid}
  </code></pre>
  <pre><code>class RoBERTaClass(torch.nn.Module):
    def __init__(self, pretrained, drop_rate, otuput_size):
      super().__init__()
      self.roberta = AutoModelForMaskedLM.from_pretrained(pretrained)
      self.drop = torch.nn.Dropout(drop_rate)
      self.fc = torch.nn.Linear(32000,otuput_size)
  
    def forward(self, ids, mask):
      out= self.roberta(ids, attention_mask=mask,return_dict=False)
      out = self.fc(self.drop(out[0]))
      return out
  </code></pre>
  <pre><code>class CreateDataset(Dataset):
    def __init__(self, X, y, tokenizer, max_len):
      self.X = X
      self.y = y
      self.tokenizer = tokenizer
      self.max_len = max_len
  
    def __len__(self):  # len(Dataset)で返す値を指定
      return len(self.y)
  
    def __getitem__(self, index):  # Dataset[index]で返す値を指定
      text = self.X[index]
      inputs = self.tokenizer.encode_plus(
        text,
        add_special_tokens=True,
        max_length=self.max_len,
        pad_to_max_length=True
      )
      ids = inputs['input_ids']
      mask = inputs['attention_mask']
  
      return {
        'ids': torch.LongTensor(ids),
        'mask': torch.LongTensor(mask),
        'labels': torch.Tensor(self.y[index])
      }
  # 最大系列長の指定
  MAX_LEN = 128
  
  # tokenizerの取得
  tokenizer = AutoTokenizer.from_pretrained(pretrained)
  #tokenizer = T5Tokenizer.from_pretrained(pretrained)
  #tokenizer.do_lower_case = True
  # Datasetの作成
  dataset_train = CreateDataset(train['text'], train[categories].values, tokenizer, MAX_LEN)
  dataset_valid = CreateDataset(valid['text'], valid[categories].values, tokenizer, MAX_LEN)
  </code></pre>
  <p>I am very new to this field. Any help and advice is very appreciated. Thanks in advance!</p>