Disabling automatic .train() for loss criteria 
❓ Questions and Help
What is your question?
I have a loss module that is loaded as part of my lightning module with its own inner network. (output is passed through the network and the result is used to compute the loss)
The problem is that when starting a train_step Lightning automatically changes the entire module to .train() which recursively changes the loss inner network as well. As the inner network has batch normalization layers this effectively makes the loss criteria change during training which is of course not desired.
Is there a right way to incorporate such a loss function?