Hi @AmanJain1008,
Thank you for posting your question in the Databricks Community.
Could you kindly check whether you are able to reproduce the issue with the below code examples:
# Import Libraries
import pandas as pd
import numpy as np
import mlflow
import mlflow.sklearn
# Load Data
data = {'Feature1': np.random.rand(10), 'Feature2':np.random.rand(10), 'Label': np.random.randint(0,2,10)}
df = pd.DataFrame(data)
# Separating features and target variable
X = df[['Feature1','Feature2']]
y = df['Label']
import tensorflow as tf
# enabling autologging for TensorFlow
mlflow.tensorflow.autolog()
# Start an MLflow run with a specified run name
with mlflow.start_run(run_name="sample_run_2") as run:
# Specify evaluations results and new process parameter values
eval_result = 4
new = 4
# Log the evaluations results and new process parameter values
mlflow.log_param("eval_result", eval_result)
mlflow.log_param("new", new)
# Train a model using TensorFlow
model = tf.keras.Sequential([
tf.keras.layers.Dense(10, input_shape=(X.shape[1],), activation='relu'),
tf.keras.layers.Dense(1, activation='sigmoid'),
])
model.compile(loss='binary_crossentropy', optimizer='adam', metrics=['accuracy'])
model.fit(X, y, epochs=10, batch_size=1)