gumbel

harvardnlp · Nov 8, 2020 · 0f3b577 · 0f3b577
1 parent 16032e0
commit 0f3b577
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torch_struct/semirings/sample.py b/torch_struct/semirings/sample.py
@@ -102,7 +102,7 @@ def forward(ctx, logits, dim):
         @staticmethod
         def backward(ctx, grad_output):
             logits, = ctx.saved_tensors
-            return logits.softmax(-1), None
+            return logits.softmax(-1) * grad_output, None
 
     class _GumbelCRFLogSumExp(torch.autograd.Function):
         @staticmethod