log all metrics on sweep run

Author: manucouto1

framework3/plugins/optimizer/wandb_optimizer.py

@@ -1,4 +1,6 @@
 from typing import Any, Dict, cast
+
+import wandb
 from framework3 import Container
 from framework3.base import BaseMetric
 from framework3.base.base_clases import BaseFilter, BasePlugin
@@ -152,8 +154,8 @@ def exec(
         match pipeline.fit(x, y):
             case None:
                 losses = pipeline.evaluate(x, y, pipeline.predict(x))
-
-                loss = losses.get(self.scorer.__class__.__name__, 0.0)
+                loss = losses.pop(self.scorer.__class__.__name__, 0.0)
+                wandb.log(dict(losses))  # type: ignore[attr-defined]
 
                 return {self.scorer.__class__.__name__: float(loss)}
             case float() as loss:

framework3/utils/wandb.py

@@ -70,6 +70,8 @@ def get_grid(aux: Dict[str, Any], config: Dict[str, Any]):
             case {"filters": filters, **r}:
                 for filter_config in filters:
                     WandbSweepManager.get_grid(filter_config, config)
+            case {"filter": filter, **r}:
+                WandbSweepManager.get_grid(filter, config)
             case {"pipeline": pipeline, **r}:  # noqa: F841
                 WandbSweepManager.get_grid(pipeline, config)
             case p_params: