diff --git a/README.md b/README.md
index 38e62800482635366b3926b6fad3e7a23eed996f..f3a9877e31f3c93134b5a7f52fad62c660cc1aa7 100644
--- a/README.md
+++ b/README.md
@@ -84,6 +84,9 @@ python scripts/generate_transfer_learning_finetune.py --input_dir <./data/exp_na
 ```
 
 
+```python
+python3 scripts/generate_transfer_learning_finetune.py --input_dir data/era5_size_64_64_3_3t_norm --dataset_hparams sequence_length=20 --checkpoint logs/era5_size_64_64_3_3t_norm/end_to_end/ours_savp --mode test --results_dir results_test_samples/era5_size_64_64_3_3t_norm/end_to_end  --batch_size 4 --dataset era5
+```
 
 ![Groud Truth](/results_test_samples/era5_size_64_64_3_norm_dup/ours_savp/Sample_Batch_id_0_Sample_1.mp4)
 # End-to-End run the entire workflow
diff --git a/scripts/generate_transfer_learning_finetune.py b/scripts/generate_transfer_learning_finetune.py
index c990ba574aba558dbd0735f85f02f82f2118ca08..9ff7fc5d255dc1f4f2666d4f9c8ac969a21528ad 100644
--- a/scripts/generate_transfer_learning_finetune.py
+++ b/scripts/generate_transfer_learning_finetune.py
@@ -4,32 +4,17 @@ from __future__ import print_function
 
 import argparse
 import errno
-import json
 import os
 import math
-import random
-import cv2
-import numpy as np
 import tensorflow as tf
-import seaborn as sns
-import pickle
-from random import seed
 import random
 import json
 import numpy as np
-#from six.moves import cPickle
 import matplotlib
 matplotlib.use('Agg')
 import matplotlib.pyplot as plt
 import matplotlib.gridspec as gridspec
-import matplotlib.animation as animation
-import seaborn as sns
-import pandas as pd
-import re
 from video_prediction import datasets, models
-from matplotlib.colors import LinearSegmentedColormap
-#from matplotlib.ticker import MaxNLocator
-#from video_prediction.utils.ffmpeg_gif import save_gif
 from skimage.metrics import structural_similarity as ssim
 import pickle
 
@@ -65,7 +50,7 @@ def main():
     parser.add_argument("--checkpoint",
                         help = "directory with checkpoint or checkpoint name (e.g. checkpoint_dir/model-200000)")
 
-    parser.add_argument("--mode", type = str, choices = ['val', 'test'], default = 'val',
+    parser.add_argument("--mode", type = str, choices = ['train','val', 'test'], default = 'val',
                         help = 'mode for dataset, val or test.')
 
     parser.add_argument("--dataset", type = str, help = "dataset class name")