Stuck at a condition with gluon.loss.Loss

[BebDong]

Version: mxnetcu101=1.5.0
OS: Ubuntu 18.04

class MyLoss(gluon.loss.Loss):
    def __init__(self, idx, threshold, **kwargs):
        super(MyLoss, self).__init__(None, 0, **kwargs)
        self._idx = idx
        self._threshold = threshold

    def hybrid_forward(self, F, logit, label):
        ...  # some operations here, let's say softmax_cross_entropy
        loss = xxxx  #  a tensor with shape (num_sample, )
        if loss[self._idx] > self._threshold:
            loss = operation 1
        else:
            loss = operation 2

        return loss

The codes run just fine on a CPU or one single GPU. But when it comes to parallel training on multiple GPUs, the program is stuck at the condition operation if loss[self._idx] > self._threshold. I have tried with F.greater as follows but it never worked. Moreover, when removing the condition, the codes run fine with either operation 1 or operation 2

def __init__:
    self._threshold = nd.array([threshold, ], dtype='float32')

def hybrid_forward:
    ...
    if F.greater(loss[self.idx], self._threshold.as_in_context(loss.context)):
        operation 1
    else:
        operation 2

[szha]

Assuming the two operations return the same shape, you can try

def __init__:
    self._threshold = self.params.get_constant('threshold', nd.array([threshold, ], dtype='float32'))

def hybrid_forward(self, F, logit, label, threshold):
    ...
    return F.where(
        F.greater(loss[self.idx], threshold),
        operation1,
        operation2)

[BebDong]

Thanks! It does work!
But what if operation 1 and operation 2 are both complicated, e.g., each contains multiple sentences, which is unrealistic to be merged?

[BebDong]

Just encapsulate complex operations as a function.