Access datamodule in custom callbacks

[vishal-wiai]

I want to create a custom Callback class where I can access certain attributes from my DataModule and log/save them before the start of the train step.

I am little confused on how to do this. Can anyone help me out with a quick snippet? Thanks!

[SkafteNicki]

Something like this should work :]

from pytorch_lightning.callbacks import Callback
class MyCallback(Callback):
    def __init__(self, ...):
        ...

    # hook for doing something with your datamodule before training step
    def on_train_batch_start(self, trainer, *args, **kwargs):
        dm = trainer.datamodule # this is a reference to your datamodule during training
        # do something here with your datamodule