[DOC] Show supported layers in error message

f-dangel · Oct 30, 2023 · 5190c9f · 5190c9f
1 parent 7161f21
commit 5190c9f
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/curvlinops/kfac.py b/curvlinops/kfac.py
@@ -321,7 +321,10 @@ def _hook_accumulate_gradient_covariance(
             covariance = einsum("bi,bj->ij", g, g).mul_(correction)
         else:
             # TODO Support convolutions
-            raise NotImplementedError(f"Layer of type {type(module)} is unsupported.")
+            raise NotImplementedError(
+                f"Layer of type {type(module)} is unsupported. "
+                + f"Supported layers: {self._SUPPORTED_MODULES}."
+            )
 
         idx = tuple(p.data_ptr() for p in module.parameters())
         if idx not in self._gradient_covariances: