diff --git a/aimodel/src/deeplabv3_plus_test_rainfall.py b/aimodel/src/deeplabv3_plus_test_rainfall.py
index 7b18483..2297215 100755
--- a/aimodel/src/deeplabv3_plus_test_rainfall.py
+++ b/aimodel/src/deeplabv3_plus_test_rainfall.py
@@ -246,14 +246,12 @@ def get_overlay(image, coloured_mask):
 
 
 def plot_samples_matplotlib(filepath, display_list, figsize=(5, 3)):
-	_, axes = plt.subplots(nrows=1, ncols=len(display_list), figsize=figsize)
-	if not isinstance(axes, list):
-		axes = [ axes ]
 	for i in range(len(display_list)):
+		plt.subplot(nrows=1, ncols=len(display_list), index=i)
 		if display_list[i].shape[-1] == 3:
-			axes[i].imshow(tf.keras.preprocessing.image.array_to_img(display_list[i]))
+			plt.imshow(tf.keras.preprocessing.image.array_to_img(display_list[i]))
 		else:
-			axes[i].imshow(display_list[i])
+			plt.imshow(display_list[i])
 	plt.savefig(filepath)