cross entropy loss with a mask

def cross_entropy_loss(inputs, labels, rescale_loss=1):
    """ cross entropy loss with a mask """
    criterion = mx.gluon.loss.SoftmaxCrossEntropyLoss(weight=rescale_loss)
    loss = criterion(inputs, labels)
    mask = S.var('mask')
    loss = loss * S.reshape(mask, shape=(-1,))
    return S.make_loss(loss.mean())