Keras / Tensorflow multi GPU InvalidArgumentError в оптимизаторе

Я хочу попробовать обучение с несколькими GPU в Keras с бэкэндом Tensorflow.

Я пытаюсь использовать функцию make_parallel, описанную здесь: https://medium.com/@kuza55/transparent-multi-gpu-training-on-tensorflow-with-keras-8b0016fd9012. Код для этого здесь (обновлен для Keras 2):

from keras.layers import concatenate
from keras.layers.core import Lambda
from keras.models import Model

import tensorflow as tf

def make_parallel(model, gpu_count):
    def get_slice(data, idx, parts):
        shape = tf.shape(data)
        size = tf.concat([ shape[:1] // parts, shape[1:] ],axis=0)
        stride = tf.concat([ shape[:1] // parts, shape[1:]*0 ],axis=0)
        start = stride * idx
        return tf.slice(data, start, size)

    outputs_all = []
    for i in range(len(model.outputs)):
        outputs_all.append([])

    #Place a copy of the model on each GPU, each getting a slice of the batch
    for i in range(gpu_count):
        with tf.device('/gpu:%d' % i):
            with tf.name_scope('tower_%d' % i) as scope:

                inputs = []
                #Slice each input into a piece for processing on this GPU
                for x in model.inputs:
                    input_shape = tuple(x.get_shape().as_list())[1:]
                    slice_n = Lambda(get_slice, output_shape=input_shape, arguments={'idx':i,'parts':gpu_count})(x)
                    inputs.append(slice_n)                

                outputs = model(inputs)

                if not isinstance(outputs, list):
                    outputs = [outputs]

                #Save all the outputs for merging back together later
                for l in range(len(outputs)):
                    outputs_all[l].append(outputs[l])

    # merge outputs on CPU
    with tf.device('/cpu:0'):
        merged = []
        for outputs in outputs_all:
            merged.append(concatenate(outputs, axis=0))

        return Model(inputs=model.inputs, outputs=merged)

Создаю модель:

model = make_parallel(create_model(...), 4)
model.compile(optimizer='adam', loss='mse', metrics=['mae', 'mse',])

После запуска Fit он тренируется в течение одной эпохи, а затем вылетает со следующим исключением:

InvalidArgumentError (see above for traceback): Incompatible shapes: [120,1] vs. [122,1]
     [[Node: training_6/Adam/gradients/loss_10/concatenate_7_loss/sub_grad/BroadcastGradientArgs = BroadcastGradientArgs[T=DT_INT32, _class=["loc:@loss_10/concatenate_7_loss/sub"], _device="/job:localhost/replica:0/task:0/gpu:0"](training_6/Adam/gradients/loss_10/concatenate_7_loss/sub_grad/Shape/_10935, training_6/Adam/gradients/loss_10/concatenate_7_loss/sub_grad/Shape_1)]]
     [[Node: training_6/Adam/gradients/concatenate_7/concat_grad/Slice_1/_11003 = _Recv[client_terminated=false, recv_device="/job:localhost/replica:0/task:0/gpu:1", send_device="/job:localhost/replica:0/task:0/cpu:0", send_device_incarnation=1, tensor_name="edge_4728_training_6/Adam/gradients/concatenate_7/concat_grad/Slice_1", tensor_type=DT_FLOAT, _device="/job:localhost/replica:0/task:0/gpu:1"]()]]

Что-то идет не так, когда доходит до стадии объединения градиентов моделей на разных графических процессорах. Несовместимые размеры формы в исключении каким-то образом связаны с размером партии (здесь 128) (т.е. изменение размера партии приводит к изменению несовместимых размеров формы).