For categorical class RuntimeError: 0D or 1D target tensor expected, multi-target not supported

I have 28 features and target variable is categorical (0-8) i.e. 9 target variable .

Data sample:

X_train.shape,y_train.shape
output --((640, 28), (640, 1))

X_train[0]
output --array([0.4546875 , 0.63958333, 0.46875   , 0.62916667, 0.4859375 ,
       0.62916667, 0.5015625 , 0.64166667, 0.4859375 , 0.65      ,
       0.4671875 , 0.65      , 0.478125  , 0.6375    , 0.5625    ,
       0.64166667, 0.5765625 , 0.62708333, 0.5921875 , 0.62708333,
       0.60625   , 0.63541667, 0.59375   , 0.64583333, 0.5765625 ,
       0.64791667, 0.58125   , 0.63541667])

y_train[0]
output --array([1])

defined data generator and model like below

class ClassifierDataset(Dataset):
    
    def __init__(self, X_data, y_data):
        self.X_data = X_data
        self.y_data = y_data
        
    def __getitem__(self, index):
        return self.X_data[index], self.y_data[index]
        
    def __len__ (self):
        return len(self.X_data)


train_dataset = ClassifierDataset(torch.from_numpy(X_train).float(), torch.from_numpy(y_train).long())
val_dataset = ClassifierDataset(torch.from_numpy(X_val).float(), torch.from_numpy(y_val).long())
test_dataset = ClassifierDataset(torch.from_numpy(X_test).float(), torch.from_numpy(y_test).long())
EPOCHS = 150
BATCH_SIZE = 32
LEARNING_RATE = 0.0007
NUM_FEATURES = len(X[0])
NUM_CLASSES = 9

train_loader = DataLoader(dataset=train_dataset,
                          batch_size=BATCH_SIZE,
                          shuffle = True
)
val_loader = DataLoader(dataset=val_dataset, batch_size=1)
test_loader = DataLoader(dataset=test_dataset, batch_size=1)


class MulticlassClassification(nn.Module):
    def __init__(self, num_feature, num_class):
        super(MulticlassClassification, self).__init__()
        
        self.layer_1 = nn.Linear(num_feature, 512)
        self.layer_2 = nn.Linear(512, 128)
        self.layer_3 = nn.Linear(128, 64)
        self.layer_out = nn.Linear(64, num_class) 
        
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(p=0.2)
        self.batchnorm1 = nn.BatchNorm1d(512)
        self.batchnorm2 = nn.BatchNorm1d(128)
        self.batchnorm3 = nn.BatchNorm1d(64)
        
    def forward(self, x):
        x = self.layer_1(x)
        x = self.batchnorm1(x)
        x = self.relu(x)
        
        x = self.layer_2(x)
        x = self.batchnorm2(x)
        x = self.relu(x)
        x = self.dropout(x)
        
        x = self.layer_3(x)
        x = self.batchnorm3(x)
        x = self.relu(x)
        x = self.dropout(x)
        
        x = self.layer_out(x)
        
        return x

Defined loss and batch size as:

model = MulticlassClassification(num_feature = NUM_FEATURES, num_class=NUM_CLASSES)
model.to(device)

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)
print(model)

Defined function for multi accuracy class

def multi_acc(y_pred, y_test):
    y_pred_softmax = torch.log_softmax(y_pred, dim = 1)
    _, y_pred_tags = torch.max(y_pred_softmax, dim = 1)    
    
    correct_pred = (y_pred_tags == y_test).float()
    acc = correct_pred.sum() / len(correct_pred)
    
    acc = torch.round(acc * 100)
    
    return acc

Started training like this

accuracy_stats = {
    'train': [],
    "val": []
}
loss_stats = {
    'train': [],
    "val": []
}
print("Begin training.")
for e in tqdm(range(1, EPOCHS+1)):
    
    # TRAINING
    train_epoch_loss = 0
    train_epoch_acc = 0
    model.train()
    for X_train_batch, y_train_batch in train_loader:
        print(X_train_batch.shape, y_train_batch.shape)
        X_train_batch, y_train_batch = X_train_batch.to(device), y_train_batch.to(device)
        optimizer.zero_grad()
        
        y_train_pred = model(X_train_batch)
        # y_train_pred = y_train_pred.unsqueeze(1)
        print(y_train_pred.shape,y_train_batch.shape)
        print(y_train_batch)
        print(y_train_pred)
        # train_loss = criterion(y_train_pred, torch.max(y_train_batch,1)[1])
        train_loss = criterion(y_train_pred, y_train_batch)
        train_acc = multi_acc(y_train_pred, y_train_batch)
        
        train_loss.backward()
        optimizer.step()
        
        train_epoch_loss += train_loss.item()
        train_epoch_acc += train_acc.item()
        
        
    # VALIDATION    
    with torch.no_grad():
        
        val_epoch_loss = 0
        val_epoch_acc = 0
        
        model.eval()
        for X_val_batch, y_val_batch in val_loader:
            X_val_batch, y_val_batch = X_val_batch.to(device), y_val_batch.to(device)
            
            y_val_pred = model(X_val_batch)
            # val_loss = criterion(y_val_pred, torch.max(y_val_batch,1)[1])       
            val_loss = criterion(y_val_pred, y_val_batch)
            val_acc = multi_acc(y_val_pred, y_val_batch)
            
            val_epoch_loss += val_loss.item()
            val_epoch_acc += val_acc.item()
    loss_stats['train'].append(train_epoch_loss/len(train_loader))
    loss_stats['val'].append(val_epoch_loss/len(val_loader))
    accuracy_stats['train'].append(train_epoch_acc/len(train_loader))
    accuracy_stats['val'].append(val_epoch_acc/len(val_loader))
                              
    
    print(f'Epoch {e+0:03}: | Train Loss: {train_epoch_loss/len(train_loader):.5f} | Val Loss: {val_epoch_loss/len(val_loader):.5f} | Train Acc: {train_epoch_acc/len(train_loader):.3f}| Val Acc: {val_epoch_acc/len(val_loader):.3f}')

This is the error I am getting:

RuntimeError                              Traceback (most recent call last)
<ipython-input-529-1d57dbd350e4> in <module>
     17         print(y_train_pred)
     18         # train_loss = criterion(y_train_pred, torch.max(y_train_batch,1)[1])
---> 19         train_loss = criterion(y_train_pred, y_train_batch)
     20         train_acc = multi_acc(y_train_pred, y_train_batch)
     21 

2 frames
/usr/local/lib/python3.8/dist-packages/torch/nn/functional.py in nll_loss(input, target, weight, size_average, ignore_index, reduce, reduction)
   2699     if size_average is not None or reduce is not None:
   2700         reduction = _Reduction.legacy_get_string(size_average, reduce)
-> 2701     return torch._C._nn.nll_loss_nd(input, target, weight, _Reduction.get_enum(reduction), ignore_index)
   2702 
   2703 

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

Any idea how to correct this? I’ve been stuck for a long time