add gpu selection

predict.py

@@ -193,9 +193,13 @@ if __name__=='__main__':
     parser.add_argument('--save_name_format', default='seg_sa_{}.nii.gz', help='save mask format. {} for ED/ES')
     parser.add_argument('--no_resample', default=False, action='store_true',help='not resample image before prediction')
     parser.add_argument('--batch_size', default=1, help='how many image slices to be sent to predictor each iteration ')
+    parser.add_argument('--gpu', default=0, help='select GPU by masking shell environment variable CUDA_VISIBLE_DEVICES')
 
     args = parser.parse_args()
 
+    ### GPU CONFIG
+    os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu)
+
     ### DATA CONFIG
     root_dir = os.path.abspath(args.root_dir)
     sequence_name=args.sequence