AzureML experiment pipeline not using CUDA with PyTorch
I am running an experiment pipeline to train my model with PyTorch and CUDA.
I created the environment as follow:
    env = Environment.from_conda_specification(model, join(model, 'conda_dependencies.yml'))
    env.docker.enabled = True
    env.environment_variables = {'MODEL_NAME': model, 'BRANCH': branch, 'COMMIT': commit}
    env.docker.base_image = DEFAULT_GPU_IMAGE

    run_config = RunConfiguration()
    run_config.environment = env
    run_config.docker = DockerConfiguration(use_docker=True)

And here is the training step:
train_step = PythonScriptStep(
      name='Model Train',
      source_directory=training_dir,
      compute_target=cluster,
      runconfig=run_config,
      script_name='train_aml.py',
      arguments=[
        '--model', model,
        '--model_output_dir', model_output_dir,
      ],
      inputs=[train_dataset.as_mount()],
      outputs=[model, model_output_dir]
    )

Even though I am using a Standard_NC12_Promo machine when I run my training script, the GPU is not picked up by PyTorch device = torch.device(&quot;cuda&quot; if torch.cuda.is_available() else &quot;cpu&quot;).
If I try running my script on the same machine but not in an experiment then the GPU is used.
Do you know any potential solutions to this?

Depending on pytorch version you might need a specific version of cuda. Try cuda 11.0.3 or cuda 11.1 from here
https://github.com/Azure/AzureML-Containers/tree/master/base/gpu
Regarding your code snippet, please move environment variables out of environment object to runconfiguration
