Использование модели с пакетной нормализацией в tf.map_fn

Question

Я реализовал модель с несколькими уровнями пакетной нормализации. Затем эта модель вызывается в реализации MAML ( https://arxiv.org/pdf/1703.03400.pdf), которая используетtensorflow.map_fn. Выполнение кода вызывает эту ошибку:TypeError: Incompatible types: <dtype: 'resource'> vs. int64. Value is 0.

Одно из возможных решений - обернуть слой нормализации партии в tf.keras.layers.TimeDistributed но поскольку я использую это для реализации MAML, т.е. разные шаги - это не временные шаги, а скорее разные задачи, я не уверен, что это правильный подход.

Вот более простой фрагмент кода (не полная реализация MAML), который показывает проблему.

import tensorflow as tf
from tensorflow import keras

mod_l1 = keras.layers.Input(shape=(10,10,3), dtype=tf.float32)
mod_l2 = keras.layers.Conv2D(
    10, (2,2), strides=2, padding='valid', data_format='channels_last', activation=None, use_bias=True
)
mod_l3 = keras.layers.BatchNormalization()
mod_l4 = keras.layers.Dense(3)

mod_out = mod_l2(mod_l1)
mod_out = mod_l3(mod_out)
mod_out = mod_l4(mod_out)
mod = keras.Model(inputs=mod_l1, outputs=mod_out)

def bla(my_in):
    in1, curr_labels = my_in

    loss = keras.losses.categorical_crossentropy(
        mod4(in1), curr_labels, from_logits=False, label_smoothing=0
    )
    return tf.reduce_mean(loss)


bla_out = tf.map_fn(
    bla,
    [
        tf.placeholder(shape=(None, None, 10, 10, 3), dtype=tf.float32),
        tf.placeholder(shape=(None, None, 3), dtype=tf.float32)
    ],
    tf.float32
)

opt = tf.train.AdamOptimizer()
loss = tf.reduce_sum(bla_out)
train = opt.minimize(loss, var_list=mod.trainable_weights)

У вас есть предложения, как избежать этого сообщения об ошибке или просто использовать TimeDistributed это правильный подход?

Спасибо.