Fix reduction of cross_entropy
Summary:
Reduction of cross_entropy is always set to "mean".
Fixed to receive the reduction.
Pull Request resolved:
@@ -30,7 +30,7 @@ def cross_entropy(input, target, *, reduction="mean", **kwargs):
"""
if target.numel() == 0 and reduction == "mean":
return input.sum() * 0.0 # connect the gradient
- return F.cross_entropy(input, target, **kwargs)
+ return F.cross_entropy(input, target, reduction=reduction, **kwargs)
class _NewEmptyTensorOp(torch.autograd.Function):
