diff --git a/aimodel/src/deeplabv3_plus_test_rainfall.py b/aimodel/src/deeplabv3_plus_test_rainfall.py
index 9973c87..7b1448c 100755
--- a/aimodel/src/deeplabv3_plus_test_rainfall.py
+++ b/aimodel/src/deeplabv3_plus_test_rainfall.py
@@ -341,7 +341,7 @@ def plot_predictions(filepath, input_items, colormap, model):
 		
 		save_samples(
 			filepath_jsonl,
-			prediction_mask
+			prediction_mask.numpy().list()
 		)
 		i += 1