tensorflow - 如何在 sagemaker 训练实例中指定 cuda 版本

Question

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

标签： tensorflowamazon-sagemaker