ModuleNotFoundError: No module named 'azureml' in ML Studio

I am learning Azure ML from Microsoft tutorials, here. The third tutorial is giving me the following error.

[stderr]Traceback (most recent call last):  
[stderr]  File "train.py", line 8, in <module>  
[stderr]    from azureml.core import Run  
[stderr]ModuleNotFoundError: No module named 'azureml'  
[stderr]  

Working with Azure ML Studio and submitting the code to the environment, I am unable to find how to resolve this error.

I have checked that the package is installed (running on Azure ML studio so this is a basic assumption, but I have tested as well). Following is the code 'run-pytorch.py' which calls the script 'train.py'

# run-pytorch.py  
from azureml.core import Workspace  
from azureml.core import Experiment  
from azureml.core import Environment  
from azureml.core import ScriptRunConfig  
  
if __name__ == "__main__":  
    ws = Workspace.from_config()  
    experiment = Experiment(workspace=ws, name='day1-experiment-train')  
    config = ScriptRunConfig(source_directory='./src',  
                             script='train.py',  
                             compute_target='cpu-cluster')  
  
    # set up pytorch environment  
    env = Environment.from_conda_specification(  
        name='pytorch-env',  
        file_path='pytorch-env.yml'  
    )  
    config.run_config.environment = env  
  
    run = experiment.submit(config)  
  
    aml_url = run.get_portal_url()  
    print(aml_url)  
    print('Success...!!!')  

The code snippet for train.py is as follows

# train.py  
import os  
import argparse  
import torch  
import torch.optim as optim  
import torchvision  
import torchvision.transforms as transforms  
from model import Net  
from azureml.core import Run  
...  
...