I’m using a GAN structure where my discriminator trains on some labeled data batches, some unlabeled data batches, and some generated data batches. I’d like to keep track_running_stats turned on, but only for the labeled data. That is, during the unlabeled and generated batches, it just uses the statistics from the labeled training so far (as it does during an eval step). Is there a way I can easily disable the updating of the batch norm statistics for these batches while still having everything else about the network in train mode? Thank you for your time!