help defining new training_step() on a callback

[malfonsoarquimea]

Hi!
I need to create a callback that once every N training steps performs the forward pass over the PL module and do some calculations.
My first approach has been to simply create a callback and then define a new training_step() within that callback that does my needed calculations. The problem is that this is calculations are not being executed.
Using the debugger I see that the callback is correctly initialized and correctly passed to the trainer, but it is not entering in this newly defined training step. Here is a minimal example of what I need to do

Do you have some insights on what I am doing wrong?

[rohitgr7]

hey @malfonsoarquimea !

Callback.training_step is not a hook so it won't be called automatically. For you use-case you can do something like:

class CustomCallback(Callback):
    def __init__(..., every_n_train_steps):
        self.every_n_train_steps = every_n_train_steps
        ...
    
    def on_train_batch_end(self, trainer, pl_module, *args, **kwargs):
        if trainer.global_step % self.every_n_train_steps == 0:
            outputs = pl_module(self.input)

[malfonsoarquimea]

Thanks very much, that is what I was looking for