GPU enabled but not used when running Transformer in TensorFlow 2.8

I am trying to run the notebook Transformer model for language understanding https://www.tensorflow.org/text/tutorials/transformer from the TensorFlow website.

Although I get the message "Plugin optimizer for device_type GPU is enabled", the GPU is not used. The code uses no more than 30% of the GPU and is slower than for the case that I use only the CPU.

When I am running a code with CNNs, it uses 100% of the GPU.

I believe this is related to a previous post of mine https://developer.apple.com/forums/thread/709142 that had LSTMs. The transformer code does not have LSTMs though.

I am using MacBook Pro 14 (10 CPU cores, 16 GPU cores) and TensorFlow-macos 2.8 with TensorFlow-metal 0.5.0.  I face the same problem for TensorFlow-macos 2.9.2 too.

My environment has:

tensorflow-macos          2.8.0  

tensorflow-metal          0.5.0 

tensorflow-text           2.8.1  

tensorflow-datasets       4.6.0                   

tensorflow-deps           2.8.0                         

tensorflow-hub            0.12.0                      

tensorflow-metadata       1.8.0   

Any idea why the GPU is not used for more than 30% although it is enabled?

GPU enabled but not used when running Transformer in TensorFlow 2.8
 
 
Q