Model Training
model = Net()
optimizer = optim.Adam(model.parameters(), lr=0.001)
loss_func = nn.BCEWithLogitsLoss()

epochs = 1
loss_list = []

model.train()

for epoch in range(epochs):
    total_loss = []
    
    for X_tr, Y_tr in train_ldr:
       X_tr = X_tr.unsqueeze(0)
       optimizer.zero_grad()
       
       output = model(X_tr)
       pred = output.argmax(dim=1, keepdim=True)
    
       loss = loss_func(output, Y_tr)
        # Backward pass
       loss.backward()
        # Optimize the weights
       optimizer.step()
        
       total_loss.append(loss.item())
       print(Y_tr, pred)
    loss_list.append(sum(total_loss)/len(total_loss))
#    print('Training [{:.0f}%]\tLoss: {:.4f}'.format(
#         100. * (epoch + 1) / epochs, loss_list[-1]))
         

This is my code which is working but with the single label class prediction error.