diff --git a/combine_images.py b/combine_images.py
new file mode 100644
index 0000000..622a207
--- /dev/null
+++ b/combine_images.py
@@ -0,0 +1,49 @@
+import argparse
+from pathlib import Path
+
+from PIL import Image
+
+
+def main():
+  parser = argparse.ArgumentParser()
+  parser.add_argument('path_to_images', type=str)
+  parser.add_argument('--output_name', type=str, default='plots.png')
+
+  args = parser.parse_args()
+
+  variables = [
+      'crossing_angle',
+      'dip_angle',
+      'drift_length',
+      'pad_coord_fraction',
+      'time_bin_fraction',
+  ]
+
+  stats = [
+      'Mean0',
+      'Mean1',
+      'Sigma0^2',
+      'Sigma1^2',
+      'Cov01',
+      'Sum',
+  ]
+
+  img_path = Path(args.path_to_images)
+  images = [[Image.open(img_path / f'{s} vs {v}_amp_gt_1.png') for v in variables] for s in stats]
+
+  width, height = images[0][0].size
+
+  new_image = Image.new('RGB', (width * len(stats), height * len(variables)))
+
+  x_offset = 0
+  for img_line in images:
+      y_offset = 0
+      for img in img_line:
+          new_image.paste(img, (x_offset, y_offset))
+          y_offset += img.size[1]
+      x_offset += img.size[0]
+
+  new_image.save(img_path / args.output_name)
+
+if __name__ == '__main__':
+  main()
diff --git a/dump_graph_model_v4.py b/dump_graph_model_v4.py
index 39f2cfa..0889bf0 100644
--- a/dump_graph_model_v4.py
+++ b/dump_graph_model_v4.py
@@ -3,20 +3,26 @@
 
 import tensorflow as tf
 
+from cuda_gpu_config import setup_gpu
 from model_export import dump_graph
-from models.baseline_v4_8x16 import preprocess_features
+from models.model_v4 import preprocess_features, Model_v4
+from models.utils import load_weights
+from run_model_v4 import load_config
 
 def main():
     parser = argparse.ArgumentParser(fromfile_prefix_chars='@')
     parser.add_argument('--checkpoint_name', type=str, required=True)
     parser.add_argument('--output_path', type=str, default='model_export/model_v4/graph.pbtxt')
     parser.add_argument('--latent_dim', type=int, default=32, required=False)
-    parser.add_argument('--dont_hack_upsampling_op', default=False, action='store_true')
+    parser.add_argument('--dont_hack_upsampling_op', default=True, action='store_true')
     parser.add_argument('--test_input', type=float, nargs=4, default=None)
     parser.add_argument('--constant_seed', type=float, default=None)
+    parser.add_argument('--gpu_num', type=str, default=None)
 
     args, _ = parser.parse_known_args()
 
+    setup_gpu(args.gpu_num)
+
     print("")
     print("----" * 10)
     print("Arguments:")
@@ -30,13 +36,10 @@ def epoch_from_name(name):
         return int(epoch)
 
     model_path = Path('saved_models') / args.checkpoint_name
-    gen_checkpoints = model_path.glob("generator_*.h5")
-    latest_gen_checkpoint = max(
-        gen_checkpoints,
-        key=lambda path: epoch_from_name(path.stem)
-    )
 
-    model = tf.keras.models.load_model(str(latest_gen_checkpoint), compile=False)
+    full_model = Model_v4(load_config(model_path / 'config.yaml'))
+    load_weights(full_model, model_path)
+    model = full_model.generator
 
     if args.constant_seed is None:
         def preprocess(x):
diff --git a/model_export/run_docker_centos.sh b/model_export/run_docker_centos.sh
index 47e1afb..4df0f0e 100755
--- a/model_export/run_docker_centos.sh
+++ b/model_export/run_docker_centos.sh
@@ -1,4 +1,5 @@
 docker run -it \
+  --rm \
   --env HOME=`pwd` \
   -v `pwd`:`pwd` \
   centos:centos7.8.2003 \
diff --git a/run_docker.sh b/run_docker.sh
index dc1fb62..2fc7e28 100755
--- a/run_docker.sh
+++ b/run_docker.sh
@@ -1 +1,8 @@
-docker run -u $(id -u):$(id -g) --runtime nvidia -v `pwd`:`pwd` silikhon/tensorflow2:v1 /bin/bash -c 'cd '`pwd`'; python test_script.py'
+docker run -it \
+  --rm \
+  -u $(id -u):$(id -g) \
+  --env HOME=`pwd` \
+  --runtime nvidia \
+  -v `pwd`:`pwd` \
+  silikhon/tensorflow2:v1 \
+  /bin/bash