pytorch update: no need to rewrap variable in backward()

78a6ef02 · Myle Ott · Sergey Edunov · 866b27d5 · 78a6ef02
Commit 78a6ef02 authored Feb 23, 2018 by Myle Ott Committed by Sergey Edunov Feb 27, 2018
Hide whitespace changes
Inline Side-by-side

Showing with 4 additions and 1 deletion

fairseq/criterions/label_smoothed_cross_entropy.py fairseq/criterions/label_smoothed_cross_entropy.py +4 -1

No files found.
--- a/fairseq/criterions/label_smoothed_cross_entropy.py
+++ b/fairseq/criterions/label_smoothed_cross_entropy.py
@@ -41,7 +41,10 @@ class LabelSmoothedNLLLoss(torch.autograd.Function):
    @staticmethod
    def backward(ctx, grad):
-        return utils.volatile_variable(ctx.grad_input) * grad, None, None, None, None, None
+        grad_input = ctx.grad_input
+        if not isinstance(grad_input, torch.autograd.Variable):
+            grad_input = utils.volatile_variable(grad_input)
+        return grad_input * grad, None, None, None, None, None
 @register_criterion('label_smoothed_cross_entropy')