diff --git a/modules/deeplearning/icing_fcn.py b/modules/deeplearning/icing_fcn.py
index f70d69d73e4e98e491c1a66f2685839c60c9f611..122a6842ee7ac38bb11fb3728764b73fa94deefd 100644
--- a/modules/deeplearning/icing_fcn.py
+++ b/modules/deeplearning/icing_fcn.py
@@ -2,7 +2,7 @@ import tensorflow as tf
 from util.setup import logdir, modeldir, cachepath, now, ancillary_path, home_dir
 from util.util import EarlyStop, normalize, make_for_full_domain_predict, get_training_parameters
 from util.geos_nav import get_navigation
-from util.augment import augment_image
+from util.augment import augment_image_3arg
 
 import os, datetime
 import numpy as np
@@ -478,7 +478,7 @@ class IcingIntensityFCN:
         dataset = dataset.cache()
         dataset = dataset.shuffle(PROC_BATCH_BUFFER_SIZE, reshuffle_each_iteration=True)
         if DO_AUGMENT:
-            dataset = dataset.map(augment_image(), num_parallel_calls=8)
+            dataset = dataset.map(augment_image_3arg(), num_parallel_calls=8)
         dataset = dataset.prefetch(buffer_size=1)
         self.train_dataset = dataset