如何在Sagemaker培训实例中指定cuda版本

Question

我目前正在使用tensorflow 2.3在AWS sagemaker上训练CNN模型。我想在GPU实例上运行训练，但似乎默认的cuda版本是cuda 10.0，但我需要的是cuda 10.1。我需要一种在训练实例中指定cuda版本的方法。我使用的是estimator.Tensorflow，它使用的是脚本文件，而不是docker映像。