fix trace computation for cases with negative curvature

amirgholami · Aug 10, 2021 · 1a42737 · 1a42737
1 parent 8a61d6c
commit 1a42737
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/pyhessian/hessian.py b/pyhessian/hessian.py
@@ -183,7 +183,7 @@ def trace(self, maxIter=100, tol=1e-3):
             else:
                 Hv = hessian_vector_product(self.gradsH, self.params, v)
             trace_vhv.append(group_product(Hv, v).cpu().item())
-            if abs(np.mean(trace_vhv) - trace) / (trace + 1e-6) < tol:
+            if abs(np.mean(trace_vhv) - trace) / (abs(trace) + 1e-6) < tol:
                 return trace_vhv
             else:
                 trace = np.mean(trace_vhv)