Fix InstructPix2Pix training in multi-GPU mode (#2978)
* fix: norm group test for UNet3D.
* fix: unet rejig.
* fix: unwrapping when running validation inputs.
* unwrapping the unet too.
* fix: device.
* better unwrapping.
* unwrapping before ema.
* unwrapping.
