Updated plot

visualize.py

@@ -60,7 +60,7 @@ def display_cifar_misclassified_data(data: list,
         img = np.transpose(img, (1, 2, 0))
         img = inv_transform(image=img)['image']
         plt.imshow(img)
-        plt.title(r"Correct: " + CLASS_NAMES[data[i][1].item()] + '\n' + 'Output: ' + CLASS_NAMES[data[i][2].item()])
+        plt.title(r"Correct: " + CLASS_NAMES[data[i][1].item()] + 'Output: ' + CLASS_NAMES[data[i][2].item()])
         plt.xticks([])
         plt.yticks([])
 
@@ -113,6 +113,6 @@ def display_gradcam_output(data: list,
 
         # Display the images on the plot
         plt.imshow(visualization)
-        plt.title(r"Correct: " + CLASS_NAMES[data[i][1].item()] + '\n' + 'Output: ' + CLASS_NAMES[data[i][2].item()])
+        plt.title(r"Correct: " + CLASS_NAMES[data[i][1].item()] + 'Output: ' + CLASS_NAMES[data[i][2].item()])
         plt.xticks([])
         plt.yticks([])