RuntimeError: 1D target tensor expected, multi-target not supported

Код модели и загрузчика данных:

class BertClassifier:

    def __init__(self, model_path, tokenizer_path, n_classes=2, epochs=1, model_save_path='/content/bert.pt'):
        self.model = BertForSequenceClassification.from_pretrained(model_path)
        #self.tokenizer = BertTokenizer.from_pretrained(tokenizer_path)
        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        self.model_save_path=model_save_path
        self.max_len = 512
        self.epochs = epochs
        self.out_features = self.model.bert.encoder.layer[1].output.dense.out_features
        self.model.i = torch.nn.Linear(self.out_features, n_classes)
        self.model.to(self.device)
    
    def preparation(self, X_train, y_train, attention_mask):#, X_valid, y_valid):
        # create datasets
        self.train_set = dataset(X_train, y_train, attention_mask)#, self.tokenizer)
        #self.valid_set = CustomDataset(X_valid, y_valid, self.tokenizer)

        # create data loaders
        self.train_loader = DataLoader(self.train_set, batch_size=1, shuffle=True)
        #self.valid_loader = DataLoader(self.valid_set, batch_size=2, shuffle=True)
        #self.X_train = X_train
        #self.y_train = y_train
        #self.attention_mask = attention_mask
        # helpers initialization
        self.optimizer = AdamW(self.model.parameters(), lr=2e-5, correct_bias=False)
        self.scheduler = get_linear_schedule_with_warmup(
                self.optimizer,
                num_warmup_steps=0,
                num_training_steps=len(self.train_loader) * self.epochs
            )
        self.loss_fn = torch.nn.CrossEntropyLoss().to(self.device)#torch.nn.BCEWithLogitsLoss#torch.nn.CrossEntropyLoss().to(self.device)
            
    def fit(self):
        self.model = self.model.train()
        losses = []
        correct_predictions = 0
        #with torch.no_grad():
        for data in self.train_loader:
            input_ids = data["input_ids"].to(self.device)
            attention_mask = data["attention_mask"].to(self.device)
            targets = data["targets"].to(self.device)

            outputs = self.model(
                    input_ids=input_ids,
                    attention_mask=attention_mask
                    )

            preds = torch.argmax(outputs.logits, dim=1)
            loss = self.loss_fn(outputs.logits, targets)

            correct_predictions += torch.sum(preds == targets)

            losses.append(loss.item())
            
            loss.backward()
            torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
            self.optimizer.step()
            self.scheduler.step()
            self.optimizer.zero_grad()

        train_acc = correct_predictions.double() / len(self.train_set)
        train_loss = np.mean(losses)
        return train_acc, train_loss
    
  #  def eval(self):
      #  self.model = self.model.eval()
      #  losses = []
      #  correct_predictions = 0

      #  with torch.no_grad():
         #   for data in self.valid_loader:
          #      input_ids = data["input_ids"].to(self.device)
          #      attention_mask = data["attention_mask"].to(self.device)
          #      targets = data["targets"].to(self.device)

         #       outputs = self.model(
          #          input_ids=input_ids,
          #          attention_mask=attention_mask
      #              )

       #         preds = torch.argmax(outputs.logits, dim=1)
       #         loss = self.loss_fn(outputs.logits, targets)
       #         correct_predictions += torch.sum(preds == targets)
        #        losses.append(loss.item())
        
       # val_acc = correct_predictions.double() / len(self.valid_set)
       # val_loss = np.mean(losses)
       # return val_acc, val_loss
    
    def train(self):
        best_accuracy = 0
        for epoch in range(self.epochs):
            print(f'Epoch {epoch + 1}/{self.epochs}')
            train_acc, train_loss = self.fit()
            print(f'Train loss {train_loss} accuracy {train_acc}')

            #val_acc, val_loss = self.eval()
            #print(f'Val loss {val_loss} accuracy {val_acc}')
            #print('-' * 10)

           # if val_acc > best_accuracy:
               # torch.save(self.model, self.model_save_path)
               # best_accuracy = val_acc

        self.model = torch.load(self.model_save_path)
    
    def predict(self, text):
        encoding = self.tokenizer.encode_plus(
            text,
            add_special_tokens=True,
            max_length=self.max_len,
            return_token_type_ids=False,
            truncation=True,
            padding='max_length',
            return_attention_mask=True,
            return_tensors='pt',
        )
        
        out = {
              'text': text,
              'input_ids': encoding['input_ids'].flatten(),
              'attention_mask': encoding['attention_mask'].flatten()
          }
        
        input_ids = out["input_ids"].to(self.device)
        attention_mask = out["attention_mask"].to(self.device)
        
        outputs = self.model(
            input_ids=input_ids.unsqueeze(0),
            attention_mask=attention_mask.unsqueeze(0)
        )
        
        prediction = torch.argmax(outputs.logits, dim=1).cpu().numpy()[0]

        return prediction

X_train:

tensor([[  0,   0,   0,  ...,   0,   0,   0],
        [  0,   0,   0,  ...,   0,   0,   0],
        [ 16,   1, 331,  ..., 326,  18,   3],
        ...,
        [325, 312, 328,  ...,   0,   0,   0],
        [  0,   0,   0,  ...,   0,   0,   0],
        [  0,   0,   0,  ...,   0,   0,   0]], dtype=torch.int

y_train:

tensor([[0, 1, 0,  ..., 0, 0, 1],
        [0, 0, 0,  ..., 0, 0, 0],
        [1, 0, 1,  ..., 0, 0, 0],
        ...,
        [0, 0, 0,  ..., 0, 0, 1],
        [0, 0, 0,  ..., 0, 0, 0],
        [0, 0, 0,  ..., 0, 0, 1]])

Код ошибки:

Epoch 1/20
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
~\AppData\Local\Temp/ipykernel_8060/2412480834.py in <module>
      7 )
      8 classifier.preparation(input_ids, y_train, attention_mask)
----> 9 classifier.train()

~\AppData\Local\Temp/ipykernel_8060/3467817459.py in train(self)
     93         for epoch in range(self.epochs):
     94             print(f'Epoch {epoch + 1}/{self.epochs}')
---> 95             train_acc, train_loss = self.fit()
     96             print(f'Train loss {train_loss} accuracy {train_acc}')
     97 

~\AppData\Local\Temp/ipykernel_8060/3467817459.py in fit(self)
     48 
     49             preds = torch.argmax(outputs.logits, dim=1)
---> 50             loss = self.loss_fn(outputs.logits, targets)
     51 
     52             correct_predictions += torch.sum(preds == targets)

~\anaconda3\envs\LikeProject\lib\site-packages\torch\nn\modules\module.py in _call_impl(self, *input, **kwargs)
   1049         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1050                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1051             return forward_call(*input, **kwargs)
   1052         # Do not call functions when jit is used
   1053         full_backward_hooks, non_full_backward_hooks = [], []

~\anaconda3\envs\LikeProject\lib\site-packages\torch\nn\modules\loss.py in forward(self, input, target)
   1118 
   1119     def forward(self, input: Tensor, target: Tensor) -> Tensor:
-> 1120         return F.cross_entropy(input, target, weight=self.weight,
   1121                                ignore_index=self.ignore_index, reduction=self.reduction)
   1122 

~\anaconda3\envs\LikeProject\lib\site-packages\torch\nn\functional.py in cross_entropy(input, target, weight, size_average, ignore_index, reduce, reduction)
   2822     if size_average is not None or reduce is not None:
   2823         reduction = _Reduction.legacy_get_string(size_average, reduce)
-> 2824     return torch._C._nn.cross_entropy_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index)
   2825 
   2826 

RuntimeError: 1D target tensor expected, multi-target not supported

Я посмотрел решение этой проблемы на разных вопросниках, понял, что нужно предоставить данные функции потерь в формате (batch_size, num_tasks, num_clasess), но не понимаю каким образом мне преобразовать y_train в этот формат.