Performance and Confusion matrix

I have one training loop for train and valid. I want the accuracy, confusion matrix for train, and validation set. I also want the matshow plot to display the numbers inside the graph. This is my code.

import numpy as np
batch_size = 32
epochs = 50
min_valid_loss = np.inf # track change in validation loss

train_losses, valid_losses = [], []
batch_train_losses, batch_valid_losses = [],[]
epoch_train_losses,epoch_valid_losses=[],[]

M = ConfusionMetrics()

for e in range(epochs):
    train_loss = 0.0
    model.train()     # Optional when not using Model Specific layer
    for data, target in train_loader:
         # Transfer Data to GPU if available
        if torch.cuda.is_available():
            data, target = data.cuda(), target.cuda()
            #print("targets.data", targets.item())
        
        
         # Clear the gradients
        optimizer.zero_grad()
        
        # Forward Pass
        output = (model(data))
        squeezed_output = torch.squeeze(output)
        
        # Find the Loss
        loss = criterion(squeezed_output.view(-1,1),target.view(-1,1))
        
        # Calculate gradients 
        loss.backward()
        # Update Weights
        optimizer.step()
        # Calculate Loss
        train_loss += loss.item()
        
        # calculate average losses
        #train_loss = train_loss/len(train_loader.dataset)
        train_losses.append(train_loss)
        epoch_train_loss =torch.tensor(train_losses).mean()
        epoch_train_losses.append(epoch_train_loss)
        
     
    #valid_loss = 0.0
    model.eval()     # Optional when not using Model Specific layer
    for data, target in valid_loader:
        if torch.cuda.is_available():
            data, target = data.cuda(), target.cuda()
        
        output = model(data)
        squeezed_output = torch.squeeze(output)
        print("=====LOSS===========",target)
        
        #loss = criterion(squeezed_output,target)
        loss = criterion(squeezed_output.view(-1,1),target.view(-1,1))
        valid_loss = loss.item() * data.size(0)
        
        
        #valid_loss = valid_loss/len(valid_loader.dataset)
        valid_losses.append(valid_loss)
        epoch_valid_loss=torch.tensor(valid_losses).mean()
        epoch_valid_losses.append(epoch_valid_loss)
        
        
        
        print('total confusion matrix')
        M.matrix()
        print('Accuracy',M.accuracy())
        print('Sensitivity', M.se())
        print('PPV', M.ppv())
        print(M.accumulate(target,output))
       

    print(f'Epoch {e+1} \t\t Training Loss: {train_loss / len(train_loader)} \t\t Validation Loss: {valid_loss / len(valid_loader)}')
    
   
    plt.matshow(M.matrix())
    reset = M.reset()
   

    if min_valid_loss > valid_loss:
        print(f'Validation Loss Decreased({min_valid_loss:.6f}--->{valid_loss:.6f}) \t Saving The Model')
        min_valid_loss = valid_loss