Skip to content

Conversation

@Atharva9621
Copy link

@Atharva9621 Atharva9621 commented Oct 20, 2025

Add support for custom loss and metrics in model_sweep

Fixes #544

  • Custom loss, metrics, and optimizers can now be passed to model_sweep in the same way as tabular_model.fit() through custom_fit_params.
  • custom_fit_params expects a dictionary specifying the custom loss, metrics, or optimizer.
  • Minimal code changes; fully backward compatible.
  • Updated corresponding tests.

Example usage

class CustomLoss(nn.Module):
      def __init__(self):
          super(CustomLoss, self).__init__()
  
      def forward(self, inputs, targets):
          loss = torch.mean((inputs - targets) ** 4)
          return 100*loss.mean()

def custom_metric(y_hat, y):
    return (y_hat - y).mean()

sweep_df, best_model = model_sweep(
    task="regression",
    train=train,
    test=val,
    data_config=data_config,
    optimizer_config=optimizer_config,
    trainer_config=trainer_config,
    model_list="lite",
    custom_fit_params = {
        "loss": CustomLoss(),
        "metrics": [custom_metric],
        "metrics_prob_inputs": [True],
        "optimizer": torch.optim.Adagrad,
    }
)

📚 Documentation preview 📚: https://pytorch-tabular--587.org.readthedocs.build/en/587/

@dosubot dosubot bot added size:M This PR changes 30-99 lines, ignoring generated files. enhancement New feature or request labels Oct 20, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request size:M This PR changes 30-99 lines, ignoring generated files.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Help: custom loss for model_sweep

1 participant