ImportError: cannot import name 'dtensor'

I used the docker now and that seems to work, but now I got a new problem. That can be seen here: https://forums.developer.nvidia.com/t/assertionerror-some-python-objects-were-not-bound-to-checkpointed-values-likely-due-to-changes-in-the-python-program/213301