Hi Diksha, Diksha I hope I can help you with this question.
For GPU optimized VMs to do machine learning tasks I would recommend to use The NC-series.
e.g. Azure NC6 series (Standard_NC6) virtual machines with K80 GPUs
Maybe this articles are also interesting for you:
https://learn.microsoft.com/en-us/azure/virtual-machines/nc-series
https://azure.microsoft.com/de-de/blog/gpus-vs-cpus-for-deployment-of-deep-learning-models/
If the reply was helpful, please don’t forget to upvote or accept it as an answer, thank you.