Horovod и TensorFlow: ошибка вычисления градиентов при использовании allgather()

Question

Я пытаюсь вычислить медиану тензора для всех пакетов и всех процессов. Однако я получил ошибкуTypeError: Expected int32, got None of type 'NoneType' instead. при звонке optimizer.compute_gradients(loss). Кажется, что вычисление градиентов плохо работает с хороводами.allgatherоперация. Вот простая иллюстрация того, чего я хотел бы достичь:

my_tensor = compute_my_tensor()  # has dimension (None, N) for example
gathered_my_tensor = hvd.allgather(my_tensor)  # this will concatenate my_tensor of all processes along first dimension
median = get_median(gathered_my_tensor)
loss = compute_loss(my_tensor, median)  # say subtract each element of my_tensor by the median, then an L2 loss
optimizer = tf.compat.v1.train.AdamOptimizer()
optimizer = hvd.DistributedOptimizer(optimizer)
gradients = optimizer.compute_gradients(loss)  # this step throws the error

Видимо allgather() операция не дифференцируема, поэтому встречается Noneзначение при вычислении градиентов? Кто-нибудь знает, как это решить?