Fix return type and tolerance in test

facebookresearch · Feb 23, 2025 · a3cbcd1 · a3cbcd1
1 parent 8411ed6
commit a3cbcd1
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/distributed_shampoo/utils/shampoo_preconditioner_list.py b/distributed_shampoo/utils/shampoo_preconditioner_list.py
@@ -1165,7 +1165,7 @@ def _adaptive_amortized_computation_frequency_criterion_below_or_equal_tolerance
         off_diagonal_summed = squared_approximate_eigenvalues.fill_diagonal_(0.0).sum()
         norm = torch.sqrt(diagonal_summed + off_diagonal_summed)
         off_diagonal_norm = torch.sqrt(off_diagonal_summed)
-        return off_diagonal_norm <= tolerance * norm
+        return bool(off_diagonal_norm <= tolerance * norm)
 
     @torch.compiler.disable
     def _amortized_computation(self) -> None:

diff --git a/distributed_shampoo/utils/tests/shampoo_preconditioner_list_test.py b/distributed_shampoo/utils/tests/shampoo_preconditioner_list_test.py
@@ -1136,7 +1136,7 @@ def test_adaptive_amortized_computation_frequency_criterion_below_tolerance(
         assert not (off_diagonal_norm <= test_tolerance1 * norm)
 
         # Below tolerance.
-        test_tolerance2 = 2e-7
+        test_tolerance2 = 3e-7
         test_criterion2 = EigenvalueCorrectedShampooPreconditionerList._adaptive_amortized_computation_frequency_criterion_below_or_equal_tolerance(
             test_factor_matrix, test_factor_matrix_eigenvectors, test_tolerance2
         )