Sagemaker mxnet container mxnet.contrib.amp import error

1

Hello!

I get the import error on
import mxnet.contrib import amp
on Sagemaker mxnet container with framework_version=‘1.4.1’, which seems to be latest version according to https://github.com/aws/sagemaker-python-sdk#mxnet-sagemaker-estimators .
I tried subprocess.run([“python”, “-m”, “pip”, “install”, “mxnet”, “–pre”, “–upgrade”]) in my training script, which gave another import error “ImportError: cannot import name ‘multi_all_finite’”

Can someone help?

Thanks!