自定义Loss前向计算时比较大小

以下的讨论是基于:
MXNet版本: mxnetcu101=1.5.0
操作系统: Ubuntu 18.04

继承mxnet.gluon.loss.Loss自定义损失函数时,涉及到一个条件运算,在多GPU上无法正常计算,不报错,一直卡在cudnn auto tune 的阶段。

def hybrid_forward(self, F, pred, label):
    loss=xxxx  #  loss 为一个 (num_sample,) 的张量 
    if loss[self.idx] > self.threshold:
        operation 1
    else:
        operation 2

self.idxself.threshold均为标量,构造损失函数的时候指定。

上述代码在CPU和单张GPU上均能正常运行,当使用多GPU并行训练的时候,就卡在语句if loss[self.idx] > self.threshold:。不使用条件运算,多GPU时仅运行operation 1或者operation 2均可。

尝试过使用如下方式,仍不能正常运行,请问是什么问题呢?

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

def hybrid_forward(self, F, pred, label):
    loss=xxxx  #  loss 为一个 (num_sample,) 的张量 
    if F.greater(loss[self.idx], self.threshold.as_in_context(loss.context)):
        operation 1
    else:
        operation 2