Как добавить горячий слой в середине модели Keras?

Question

У меня есть модель автоэнкодера, подобная этой

    input_img = Input(shape=(132,))

    l1 = Dense(64, activation='relu')(input_img)
    l2 = Dense(32, activation='relu')(l1)
    l3 = Dense(nb_class)(l2)
    l_select = Activation("softmax")(l3)        
    l4 = Dense(32, activation='relu')(l_select)
    l5 = Dense(64, activation='relu')(l4)
    l6 = Dense(132, activation='relu')(l5)
    self.model_3 = Model(inputs=[input_img], output=l6)

, и я хочу вставить один горячий слой после слоя softmax, такой как input_img = Input (shape = (132,))

    l1 = Dense(64, activation='relu')(input_img)
    l2 = Dense(32, activation='relu')(l1)
    l3 = Dense(nb_class)(l2)
    l_select = Activation("softmax")(l3)
    l7 = Lambda(K.round)(l_select)

    l4 = Dense(32, activation='relu')(l7)
    l5 = Dense(64, activation='relu')(l4)
    l6 = Dense(132, activation='relu')(l5)
    self.model_2 = Model(inputs=[input_img], output=l7)
    self.model_3 = Model(inputs=[input_img], output=l6)

Я не хочу, чтобы какой-либо слой попал в мой слой декодирования, я хочу определенный выбор в слое декодирования.

Если я попытаюсь использовать K.argmax, K.round, он выведетошибка

     ValueError: An operation has `None` for gradient. Please make sure that 
     all of your ops have a gradient defined 
     (i.e. are differentiable). Common ops without gradient: K.argmax, 
     K.round, K.eval.

Есть ли способ вставить горячий слой в середину модели?