No registered 'MaxPoolWithArgmax' OpKernel for XLA_GPU_JIT devices compatible with node

TensorFlow General Discussion
1

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?

3

Hi @Dmitry_Redgrave, Could please try by training the model with the below code.

with tf.device('/device:GPU:0'):
    history=model.fit()

Thank You.