Facing this error: Error running graph component for node train_DIETClassifier

I built a rasa chatbot using rasa init in my local machine and when following the similar process on my Virtual Machine(Azure), I am facing the following error of DIETClassifier. Is it becasue of GPU that’s available on my VM? Please help this is urgent for me