diff --git a/train_scripts/create_dataset_overlay.py b/train_scripts/create_dataset_overlay.py
index c469871..6e44b1c 100755
--- a/train_scripts/create_dataset_overlay.py
+++ b/train_scripts/create_dataset_overlay.py
@@ -107,7 +107,7 @@ def to_uint8(channel):
 
 
 def save_notebook_style_png(path, channel):
-    fig = plt.figure()
+    fig = plt.figure(figsize=(16, 16))
     plt.imshow(channel)
     plt.savefig(path)
     plt.clf()