diff --git a/BrainMaGe/tester/test_single_run.py b/BrainMaGe/tester/test_single_run.py
index eedd398..ac11d74 100755
--- a/BrainMaGe/tester/test_single_run.py
+++ b/BrainMaGe/tester/test_single_run.py
@@ -81,8 +81,10 @@ def infer_single_ma(input_path, output_path, weights, mask_path=None, device="cp
 
     if mask_path is not None:
         print("You chose to save the brain. We are now saving it with the masks.")
-        image_data[to_save == 0] = 0
-        to_save_brain = nib.Nifti1Image(image_data, patient_nib.affine)
+        patient_nib_write = nib.load(input_path)
+        image_data_write = patient_nib_write.get_fdata()
+        image_data_write[to_save == 0] = 0
+        to_save_brain = nib.Nifti1Image(image_data_write, patient_nib_write.affine)
         nib.save(to_save_brain, os.path.join(mask_path))
 
     print("Thank you for using BrainMaGe")