When the generator is trying run a NetworkPolicy it has a batch size of 1 which doesn't work well with BatchNorm Layers (see error in title). The error is due to the fact that the mean and var of the BatchNorm are still getting updated even with torch.no_grad and the error for batch size 1 is probably to indicate that this is not the intended behavior.
Do you think it is okay to run the NetworkPolicy in eval mode?
Good point, I agree the network should be run in eval mode in NetworkPolicy! I probably missed it, because I did not use layers that need eval setting during inference.
When the generator is trying run a NetworkPolicy it has a batch size of 1 which doesn't work well with BatchNorm Layers (see error in title). The error is due to the fact that the mean and var of the BatchNorm are still getting updated even with torch.no_grad and the error for batch size 1 is probably to indicate that this is not the intended behavior.
Do you think it is okay to run the NetworkPolicy in eval mode?