using gpu memory but not computation

[psm1000]

I use A6000 gpu installed with nvidia driver version 11.6
but I use nvidia/cuda:11.2.0. image where cudnn version is 8.1.1

I installed all the packages in requirements.txt

The training is proceeding but GPU uses only memory and doesn't compute. I think that computation might be proceeded in the cpu.

Also, I change 'train.py' like this,

def main(argv):
jax.config.parse_flags_with_absl()
physical_devices = tf.config.list_physical_devices('GPU')
tf.config.experimental.set_visible_devices(physical_devices[2:], 'GPU')

So I uninstalled tensorflow and newly install tensorflow-gpu, but the result is same.

I need help.