This commit is contained in:
Jeremy Hummel 2019-08-14 11:06:38 -07:00
commit f3ef017ff4

View file

@ -436,7 +436,8 @@ class SAEModel(ModelBase):
else: else:
if self.options['learn_mask']: if self.options['learn_mask']:
self.AE_convert = K.function([warped_dst], [pred_src_dst[-1], pred_dst_dstm[-1], pred_src_dstm[-1]]) 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: else:
self.AE_convert = K.function([warped_dst], [pred_src_dst[-1]]) self.AE_convert = K.function([warped_dst], [pred_src_dst[-1]])