How to save additional variables while checkpointing?

[arunpatro-meesho]

In vanilla pytorch, I save like this:

if v_loss < best_val_loss:
    print("Found better model, saving")
    model_save_path = f"models/{sscat_id}/{attribute}/ts={time_stamp}/best.pth"
    best_val_loss = v_loss
    torch.save(
        {
            "epoch": epoch,
            "model_state_dict": model.state_dict(),
            "optim_state_dict": optimizer.state_dict(),
            "report": report,
            "label_encoder_dict": label_encoder_dict,
            "inverse_label_encoder_dict": {v: k for k, v in label_encoder_dict.items()},
            "weighted_f1": weighted_f1,
        },
        model_save_path,

How can I save extra keys like report, label_encoder_dict etc while using ModelCheckpoint callbacks? How do I save a variable inside the LightningModule instance?

[rohitgr7]

you can use on_save_checkpoint hook inside LightningModule.