enable gradient calculation in dist.topk() by default (#56)

harvardnlp · Mar 10, 2020 · d921f30 · d921f30
1 parent 0e47fe9
commit d921f30
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/torch_struct/distributions.py b/torch_struct/distributions.py
@@ -93,9 +93,10 @@ def topk(self, k):
         Returns:
             kmax (*k x batch_shape x event_shape*)
         """
-        return self._struct(KMaxSemiring(k)).marginals(
-            self.log_potentials, self.lengths, _raw=True
-        )
+        with torch.enable_grad():
+            return self._struct(KMaxSemiring(k)).marginals(
+                self.log_potentials, self.lengths, _raw=True
+            )
 
     @lazy_property
     def mode(self):