diff --git a/Enhancement_GAN.py b/Enhancement_GAN.py
index 6f3a971..d16e37c 100644
--- a/Enhancement_GAN.py
+++ b/Enhancement_GAN.py
@@ -1,3 +1,8 @@
+#!/usr/bin/python3
+
+
+
+#Libraries that we used in this python script
 import os
 import time
 import numpy as np
@@ -7,6 +12,7 @@
 import sys
 import nibabel
 import SimpleITK as sitk
+import sys
 from keras.models import model_from_json
 
 # load GAN model here
@@ -297,14 +303,18 @@ def enhance(config_file):
         print(volume_name_enhanced)
     
     #print(data_num)
-
-if __name__ == '__main__':
+ 
+def main():
     if(len(sys.argv) != 2):
         print('Number of arguments should be 2. e.g.')
         print('python Enhancement_GAN.py config17/train_wt_ax.txt')
-        exit()
+        sys.exit() #to safely exit & cleanup memory
     config_file = str(sys.argv[1])
     assert(os.path.isfile(config_file))
     enhance(config_file)
     
+if __name__ == '__main__':main()
+    
+#END of script
+