Lightning not setting my model into training mode with training_step implemented

[AndreCorreaSantos]

I've implemented a model trainer with Lightning, with all the required methods, such as training_step, validation_step, and configure_optimizers. Both training_step and validation_step call my model’s forward pass, return the loss, and log the model’s performance every 10 steps on rank 0. This is all standard. But for some reason, the model is not being automatically set to training mode during training—it’s always stuck in eval mode. Is there something else I need to configure to make sure the model is properly recognized? The forward pass and everything else work fine; the issue is strictly that it isn’t being set to training mode automatically, and from the training curves it looks like the weights aren’t being updated.

To provide context, I'm fitting a model which contains several submodules and it seems some of them are not being set to training during training but others are. Is there somewhere I need to "register" nested models so that lightning trains them automatically?

Thanks in advance and any materials or sources would be greatly appreciated.

[deependujha]

Hi @AndreCorreaSantos,

could you share a minimal reproducible example of the issue? Ideally in Lightning Studio, so we can run it directly, verify the behavior, and help you resolve it more effectively.