More lenient in accepting errors of different types

FalkonML · May 1, 2024 · 51d0885 · 51d0885
1 parent fe672ec
commit 51d0885
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/falkon/models/falkon.py b/falkon/models/falkon.py
@@ -1,7 +1,7 @@
 import dataclasses
 import time
 import warnings
-from typing import Callable, Optional, Tuple, Union
+from typing import Callable, Optional, Tuple, Union, Any
 
 import torch
 from torch import Tensor
@@ -125,7 +125,7 @@ def __init__(
         center_selection: Union[str, falkon.center_selection.CenterSelector] = "uniform",
         maxiter: int = 20,
         seed: Optional[int] = None,
-        error_fn: Optional[Callable[[torch.Tensor, torch.Tensor], Union[float, Tuple[float, str]]]] = None,
+        error_fn: Optional[Callable[[torch.Tensor, torch.Tensor], Union[Any, Tuple[Any, str]]]] = None,
         error_every: Optional[int] = 1,
         weight_fn: Optional[Callable[[torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor]] = None,
         options: Optional[FalkonOptions] = None,

diff --git a/falkon/models/model_utils.py b/falkon/models/model_utils.py
@@ -109,7 +109,7 @@ def val_cback(it, beta, train_time):
             if isinstance(err, tuple) and len(err) == 2:
                 err, err_name = err
             print(
-                f"Iteration {it:3d} - Elapsed {self.fit_times_[-1]:.2f}s - {err_str} {err_name}: {err:.8f}",
+                f"Iteration {it:3d} - Elapsed {self.fit_times_[-1]:.2f}s - {err_str} {err_name}: {str(err)}",
                 flush=True,
             )
             self.val_errors_.append(err)