diff --git a/models/Model_SAE/Model.py b/models/Model_SAE/Model.py
index 4044ee9..eb6ef79 100644
--- a/models/Model_SAE/Model.py
+++ b/models/Model_SAE/Model.py
@@ -436,10 +436,9 @@ class SAEModel(ModelBase):
 
         else:
             if self.options['learn_mask']:
-                from keras import K
                 self.AE_convert = K.function([warped_dst], [pred_src_dst[-1][:4], pred_dst_dstm[-1], pred_src_dstm[-1]])
             else:
-                self.AE_convert = K.function([warped_dst], [pred_src_dst[-1]])
+                self.AE_convert = K.function([warped_dst], [pred_src_dst[-1][:4]])
 
         if self.is_training_mode:
             self.src_sample_losses = []