How to train Custom Tensorflow Models in Azure ML Studio Designer

I am currently trying out different architectures with Azure ML Ecosystem. Currently, I am testing out Azure ML Studio Designer.

When I created a custom Tensorflow model using the "Create Python Model" Component. When I run the designer pipeline I get an error saying that Tensorlfow is not found.

Error:

---------- Start of error message from Python interpreter ----------
Got exception when importing script: 'No module named 'tensorflow''.
---------- End of error message from Python  interpreter  ----------

Script in the Task:

# The script MUST define a class named AzureMLModel.
# This class MUST at least define the following three methods: "__init__", "train" and "predict".
# The signatures (method and argument names) of all these methods MUST be exactly the same as the following example.

# Please do not install extra packages such as "pip install xgboost" in this script,
# otherwise errors will be raised when reading models in down-stream modules.

import pandas as pd
import numpy as np

import tensorflow as tf
from sklearn.preprocessing import OneHotEncoder

class AzureMLModel:
    # The __init__ method is only invoked in module "Create Python Model",
    # and will not be invoked again in the following modules "Train Model" and "Score Model".
    # The attributes defined in the __init__ method are preserved and usable in the train and predict method.
    def __init__(self):
        # self.model must be assigned
        model = tf.keras.Sequential()

        model.add(tf.keras.layers.Convolution1D(filters=2, kernel_size=1,input_shape=(4,1), name='Conv1'))
        model.add(tf.keras.layers.Flatten())
        model.add(tf.keras.layers.Dense(10, activation='relu', name='Dense1'))
        model.add(tf.keras.layers.Dense(10, activation='relu', name='Dense2'))
        model.add(tf.keras.layers.Dense(3, activation='softmax', name='output'))

        optimizer = tf.keras.optimizers.Adam(lr=0.001)
        model.compile(optimizer, loss='categorical_crossentropy', metrics=['accuracy'])

        self.model = model
        self.feature_column_names = list()

    # Train model
    #   Param<df_train>: a pandas.DataFrame
    #   Param<df_label>: a pandas.Series
    def train(self, df_train, df_label):
        # self.feature_column_names records the column names used for training.
        # It is recommended to set this attribute before training so that the
        # feature columns used in predict and train methods have the same names.
        self.feature_column_names = df_train.columns.tolist()
        encoder = OneHotEncoder(sparse=False)
        df_label = encoder.fit_transform(df_label)
        
        ES = tf.keras.callbacks.EarlyStopping(monitor="val_loss", patience=10)
        
        self.model.fit(df_train, df_label, validation_split=0.1 ,epochs=1000, callbacks=[ES])

    # Predict results
    #   Param<df>: a pandas.DataFrame
    #   Must return a pandas.DataFrame
    def predict(self, df):
        # The feature columns used for prediction MUST have the same names as the ones for training.
        # The name of score column ("Scored Labels" in this case) MUST be different from any other
        # columns in input data.
        pred = self.model.predict(df[self.feature_column_names])
        return pd.DataFrame({'Scored Labels': np.argmax(pred, axis=1)})

How can I resolve this? I tried the model in a notebook and worked so there is no syntax error, just the issue with Tensorflow.