How to train JEM with batch norm

[scilearner]

Hi

Thanks for the exciting work and open-sourced code.

I have a problem. As you noted in the paper "we have been able to successfully train Joint-EBMs with Batch Normalization", I run your code with enabling batch norm --norm batch, but it seems the code doesn't work at all.

Do I miss anything? Do I need to change some code to enable batch norm?

Any help is appreciated. Thank you.