how to mlflow autolog with custom parameters

I'm trying to log my ML trials with mlflow.keras.autolog and mlflow.log_param simultaneously (mlflow v 1.22.0). However, the only things that are recorded are autolog's products, but not those of log_param.

experiment = mlf_client.get_experiment_by_name(experiment_name)
with mlflow.start_run(experiment_id=experiment.experiment_id):
    mlflow.keras.autolog(log_input_examples=True)
    mlflow.log_param('batch_size', self.batch_size)
    mlflow.log_param('training_set_size', len(kwargs['training_ID_list']))
    mlflow.log_param('testing_set_size', len(kwargs['testing_ID_list']))
    
    history = self.train_NN_model(**kwargs)

I know I can use log_param with log_model to save the model itself, but then I lose some useful stuff that autolog can record for me automatically (e.g., model summary).

Is it possible to use autolog with custom parameters for logging?