import tensorflow as tf

class Model(tf.keras.Model):
    def call(self, x):
        y = x / 1000
        for i in range(3):
            y = tf.matmul(y, x / 1000)
        return tf.reduce_sum(y, axis=[1, 2])

def get_dataset(ictx):
    ds = tf.data.Dataset.range(1, 1001, output_type=tf.float32)
    ds = ds.map(lambda i: (tf.ones((1024 * 5, 1024 * 5)) / i, 0.0))
    ds = ds.batch(8)
    ds = ds.apply(tf.data.experimental.prefetch_to_device('gpu'))
    return ds

strategy = tf.distribute.MirroredStrategy()
with strategy.scope():
    ds = strategy.distribute_datasets_from_function(get_dataset)
    model = Model()
    model.compile(loss='mse')
    model.fit(
        ds,
        epochs=1,
        steps_per_epoch=30,
        callbacks=tf.keras.callbacks.TensorBoard(profile_batch=(15, 25)))