AMP - AttributeError: module 'tensorflow.python.training.training' has no attribute 'experimental'

Hi!
I’m running NGC TF container version 18.03-py3, and getting the following error when trying to wrap-up optimizer for AMP approximations:

AttributeError: module 'tensorflow.python.training.training' has no attribute 'experimental'

Code snippet:

opt = tf.train.AdamOptimizer(learning_rate=learning_rate_tensor, name=name)
opt = tf.train.experimental.enable_mixed_precision_graph_rewrite(opt)

To import libraries, I’m using only

import TensorFlow as tf

What am I missing?

18.03 NGC containers did not support AMP. Please update to 19.07 or newer.

Hi @nleuhr.
Thanks for the update, I’ve updated the container and now AMP is working.