Correct Usage of PyTorch Lightning + Hydra + AzureML

Hey, I was wondering if there is a correct way to integrate PyTorch Lightning training scripts within AzureML? I’m currently using PyTorch Lightning to train a simple image classifier where configs are passed to the training script using Hydra. This approach works fine on my local machine which is running with x1 GPU. However, I’m experimenting running the same training script within AzureML, potentially where I could use multi-GPU compute.

Traditionally, when training on AzureML we can just use an entry script for training. Would the correct way to train using AzureML be to pass the training script as an entry script into AzureML:

image

Where I have specified N GPU which matches the GPU compute that I am using? It feels like it should work, but I’m not sure if its the optimal way of utilising multi-GPU compute within Azure/AzureML

Thanks for any suggestions!