Bug fix for saving and loading custom loss functions (#415)

* bug fix for custom loss weight state dict error

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fixed custom loss, metric setting in load model

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: manujosephv

src/pytorch_tabular/tabular_model.py

@@ -459,21 +459,40 @@ def load_model(cls, dir: str, map_location=None, strict=True):
             model_args["optimizer_params"] = {}  # For compatibility. Not Used
 
         # Initializing with default metrics, losses, and optimizers. Will revert once initialized
-        model = model_callable.load_from_checkpoint(
-            checkpoint_path=os.path.join(dir, "model.ckpt"),
-            map_location=map_location,
-            strict=strict,
-            **model_args,
-        )
-        # Updating config with custom parameters for experiment tracking
-        if custom_params.get("custom_loss") is not None:
-            model.custom_loss = custom_params["custom_loss"]
-        if custom_params.get("custom_metrics") is not None:
-            model.custom_metrics = custom_params["custom_metrics"]
+        try:
+            model = model_callable.load_from_checkpoint(
+                checkpoint_path=os.path.join(dir, "model.ckpt"),
+                map_location=map_location,
+                strict=strict,
+                **model_args,
+            )
+        except RuntimeError as e:
+            if (
+                "Unexpected key(s) in state_dict" in str(e)
+                and "loss.weight" in str(e)
+                and "custom_loss.weight" in str(e)
+            ):
+                # Custom loss will be loaded after the model is initialized
+                # continuing with strict=False
+                model = model_callable.load_from_checkpoint(
+                    checkpoint_path=os.path.join(dir, "model.ckpt"),
+                    map_location=map_location,
+                    strict=False,
+                    **model_args,
+                )
+            else:
+                raise e
         if custom_params.get("custom_optimizer") is not None:
             model.custom_optimizer = custom_params["custom_optimizer"]
         if custom_params.get("custom_optimizer_params") is not None:
             model.custom_optimizer_params = custom_params["custom_optimizer_params"]
+        if custom_params.get("custom_loss") is not None:
+            model.loss = custom_params["custom_loss"]
+        if custom_params.get("custom_metrics") is not None:
+            model.custom_metrics = custom_params.get("custom_metrics")
+            model.hparams.metrics = [m.__name__ for m in custom_params.get("custom_metrics")]
+            model.hparams.metrics_params = [{}]
+            model.hparams.metrics_prob_input = custom_params.get("custom_metrics_prob_inputs")
         model._setup_loss()
         model._setup_metrics()
         tabular_model = cls(config=config, model_callable=model_callable)