set state in LightningModule using LightningDataModule after setup

[abefrandsen]

I would like to be able to set an attribute on my LightningModule model object after the trainer has called setup(stage='fit') on the datamodule, but before the sanity validation loop and training loops begin (the reason is this: the datamodule prepares a certain tokenizer in its prepare_data method, and then loads this tokenizer in its setup method, and my model needs access to the tokenizer). I searched through the model hooks and callbacks and it's not clear if there's a hook that gets called at some point between the datamodule setup and the start of the validation/fit loops. How would I go about doing this?

[abefrandsen]

Ok I found a way to do it. I have something like the following in the datamodule. This is functional but I'd be interested if there a better pattern I can employ here.

class MyDM(L.LightningDataModule):
    ...
    def setup(self, stage):
        self._load_tokenizer()
        if hasattr(self, "trainer"): # here I communicate the tokenizer to the model
            self.trainer.lightning_module.tokenizer= self.tokenizer