Mode collapse

[muxgt]

Hi! I'm trying to train your model with tpu. Had to remove MirroredStrategy and adapt code a bit. But it somehow degrades to a mode collapse on 256-512 resolution. Like this: There are 8 workers on tpu so it should be the same as 8 gpus. I used the same batch sizes and learning rates as in the official code but I see that you have smaller batch sizes with the same learning rates. What is a correct way to match batch sizes and learning rate?