Apply method of a module

[GoktugGuvercin]

Hello Flax community,

I have a resnet architecture with batch-norm layers. While doing forward pass by resnet.apply(), I am confused at some point:
Documentation specifies that the argument batch_stats needs to be marked as mutable. Is this valid for every case or just train=True ?

For example, I am in testing phase, and I want to check performance of my model with specific metrics. In that case, I will use my test data, and batch norm layers will use its statistics. So, batch_stats will not be updated right; I think no need for mutable argument in testing phase ? What do you think ?

Thanks in advance.

[cgarciae]

Hey @GoktugGuvercin, you only need to specify mutable=True when use_running_average=False. If you are using the usual train argument, you don't need to specify mutable when train=False i.e. during inference.

[GoktugGuvercin]

Ah okey. Thank you so much for the clarification :)