Problem in U-Net implementation

help me plz @the-dharma-bum

with torch.no_grad():
        output1 = model((train_x.float()))     

        softmax = torch.exp(output1).cpu()    
        prob1 = list(softmax.numpy())
        predictions = np.argmax(prob1, axis=1)
        print('Validation accuracy train: {:.4f}%'.format(float(accuracy_score(train_y, predictions, normalize =False) * 100)))