It does say NaN loss in the logs, but only with apex/amp. Training works fine without apex/amp. If I add torch.backends.cudnn.benchmark = False before training, it gives me a Zero Division error instead of the memory access error. Any idea what I can do? The pix2pix docs say it works.