Custom training Tensorflow on Sagemaker: WARNING keras.py:603 cannot save layer values during forward pass with tf.function

Question

I am training a custom model with a custom layer using Tensorflow on Sagemaker. It runs smoothly local, no errors, no warnings. When I run it on Sagemaker, it produces many warn