I defined MaxPoolWithArgmax layer that uses tf.nn.max_pool_with_argmax function. When training on the CPU there are no problems, but on the GPU one an error occurs (specified in the title). Full output: 2024-05-01 13:24:06.597101: E external/local_xla/xla/stream_executor/cuda/cuda_d - Pastebin.com
How to fix it?