… fix it'll crash when training on GPU, because for GPU, format changes (NHWC to NCHW) are done in order to speed up all conv operations. We are doing preprocessing on CPU. But depthwise conv function is executing as if it's on GPU. Hence it's giving format error `Depthwise convolution on CPU is only supported for NHWC format`. PiperOrigin-RevId: 344328724