Is your feature request related to a problem? Please describe.
Thanks for @deepib 's asking, I will try to update the base docker and CI tests to PyTorch 22.03:
https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch
And then run our integration tests to verify it.