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

MXNet Gluon
#1

以下的讨论是基于:
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.idx 和 self.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