@KumaranT I did it already with the same result

import mlflow.pyfunc

# Load the model as a PyFuncModel
model = mlflow.pyfunc.load_model(model_uri=f"{model_version_uri}")

# Create a Spark UDF for scoring
predict_udf = mlflow.pyfunc.spark_udf(spark, model_uri=f"{model_version_uri}")

#Apply the UDF to the DataFrame
predictions_df = df_train.withColumn("predictions", predict_udf(*df_train.columns))