Best practice for saving only specific parts of the pl.LightningModule

[DerEchteFeuerpfeil]

Hello community, I have a question regarding best practice.

I am using TorchMetrics to track CLIP score and FID score in my VAE training. For that, I initialized these metrics in my VAEXperiment like so:

class VAEXperiment(pl.LightningModule):

    def __init__(self,
                 vae_model: BaseVAE,
                 params: dict) -> None:
        super(VAEXperiment, self).__init__()

        self.model = vae_model
        self.params = params
        self.curr_device = None
        self.hold_graph = False
        
        self.fid = FrechetInceptionDistance(feature=768, reset_real_features=False, normalize=True)
        self.clip_sim = CLIPScore(model_name_or_path=self.params["clip_sim_model"])

I am saving these to the LightningModule because I need to access them at various points in training and validation without reinitialising them, e.g. in on_validation_end. The only model I am interested in training and saving is the VAE (self.model). The rest is just for tracking metrics.

When I start training with the Trainer however, I see the following in my shell:

  | Name     | Type                     | Params
------------------------------------------------------
0 | model    | VanillaVAE               | 3.9 M
1 | fid      | FrechetInceptionDistance | 23.9 M
2 | clip_sim | CLIPScore                | 427 M
------------------------------------------------------
431 M     Trainable params
23.9 M    Non-trainable params
455 M     Total params
1,821.620 Total estimated model params size (MB)

which makes sense because I added them to the LightningModule directly. My issue now is, I don't want to save these, they just bloat my checkpoint files. But without code modification, they get saved automatically.

What is the best practice here? Is TorchMetrics not supposed to be in the LightningModule? Can I remove FID and CLIP from the checkpoint in on_save_checkpoint? If so, is this regarded as best practice?

Looking forward to your feedback :)