.

harvardnlp · Nov 2, 2019 · 01da3ed · 01da3ed
1 parent 92617d9
commit 01da3ed
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/torch_struct/alignment.py b/torch_struct/alignment.py
@@ -2,6 +2,7 @@
 from .helpers import _Struct
 import math
 from .sparse import *
+from pytorch_memlab import MemReporter
 
 class Alignment(_Struct):
     def _check_potentials(self, edge, lengths=None):
@@ -147,6 +148,9 @@ def merge(x, size, rsize):
             size = int(size / 2)
             rsize = rsize * 2
             chart[n] = merge(chart[n - 1], size, rsize+1)
+        reporter = MemReporter()
+        reporter.report()
+
         v = chart[-1][:, :, 0, M, N, Mid]
         return v, [log_potentials], None