Skip to content

Commit

Permalink
Fix docstring and error catching in test
Browse files Browse the repository at this point in the history
  • Loading branch information
runame committed Sep 21, 2024
1 parent cf45ef5 commit 7422b00
Show file tree
Hide file tree
Showing 2 changed files with 5 additions and 1 deletion.
4 changes: 4 additions & 0 deletions curvlinops/kfac.py
Original file line number Diff line number Diff line change
Expand Up @@ -243,7 +243,11 @@ def __init__(
Raises:
RuntimeError: If the check for deterministic behavior fails.
ValueError: If the loss function is not supported.
ValueError: If the Fisher type is not supported.
ValueError: If the KFAC approximation type is not supported.
ValueError: If ``fisher_type != FisherType.MC`` and ``mc_samples != 1``.
NotImplementedError: If ``correct_eigenvalues`` and ``fisher_type ==
FisherType.FORWARD_ONLY``.
ValueError: If ``X`` is not a tensor and ``batch_size_fn`` is not specified.
"""
if not isinstance(loss_func, self._SUPPORTED_LOSSES):
Expand Down
2 changes: 1 addition & 1 deletion test/test_kfac.py
Original file line number Diff line number Diff line change
Expand Up @@ -512,7 +512,7 @@ def test_multi_dim_output(
# KFAC for deep linear network with 4d input and output
params = list(model.parameters())
context = (
raises(ValueError, match="eigenvalues")
raises(NotImplementedError, match="eigenvalues")
if correct_eigenvalues and fisher_type == FisherType.FORWARD_ONLY
else nullcontext()
) # EKFAC for FOOF is currently not supported
Expand Down

0 comments on commit 7422b00

Please sign in to comment.