code: def train_model(**fit_kwargs):
      model = paddle.Model(LinearRegression())
      optim = paddle.optimizer.Adam(
          learning_rate=0.01, parameters=model.parameters()
      )
      model.prepare(optim, paddle.nn.MSELoss())
      train_dataset, eval_dataset = get_datasets()
      model.fit(
          train_dataset,
          eval_dataset,
          batch_size=16,
          epochs=NUM_EPOCHS,
          verbose=1,
          **fit_kwargs,
      )
      return model
  
file_path: mlflow/tests/paddle/test_paddle_autolog.py
imports: import paddle
  import pytest
  import mlflow
  from mlflow import MlflowClient
label_desc: Use image node_id 2 for calling a global function with example usage: train_model() and returns: model