pip install arcus-azureml missing azureml-sdk automatic import

[ALaks96]

Describe the bug
When trying to launch training script (set up with arcus), following error pops up:
ModuleNotFoundError: No module named 'azureml.train'

To Reproduce
Create AZureMLEnvironment, start an expirment, setup train, and launch training:

from arcus.azureml.environment.aml_environment import AzureMLEnvironment

work_env = AzureMLEnvironment.Create(config_file="../.azureml/config.json")

training_name = 'your_training_name'
trainer = work_env.start_experiment(training_name)
trainer.setup_training(training_name, overwrite=False)

dataset_name = 'your_dataset_name'

arguments = {
    '--epochs': 75,
    '--batch_size': 256,
    '--es_patience': 20,
    '--train_test_split_ratio': 0.08
}
trainer.start_training(training_name, estimator_type='tensorflow', 
                       input_datasets = [dataset_name], 
                       compute_target='your_instance', gpu_compute=True, script_parameters = arguments)

Expected behavior
Should launch new training in previously created expirement, using parameters passed to arcus in model fitting

Additional context
Installing manually the azureml-sdk package seems to solve the problem:
pip3 install azureml-sdk