[Python] mlflow unable to generate input output schema when using log_model()

I am using the latest mlflow 3.3.2 with a simple PyTorch implementation of SRCNN. I am able to log the parameters as well as the metrics for the training process as well as generate custom visualizations for the losses and PSNR alongside the model with all auxilliary files required for exporting as artifacts.

I would like to add the input output schema to each version of the model automatically. According to this document all I need to do is use log_model() with input_example as a parameter, which holds the value of a sample input.

train_dataset = TrainDataset(args.train_file)
train_dataloader = DataLoader(dataset=train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True,
drop_last=True)
eval_dataset = EvalDataset(args.eval_file)
eval_dataloader = DataLoader(dataset=eval_dataset, batch_size=1)

best_weights = copy.deepcopy(model.state_dict())
best_epoch = 0
best_psnr = 0.0

# signature = infer_signature(inputs.numpy(), model(inputs).detach().numpy())
dataiter = iter(train_dataloader)
inputs_example, labels_example = next(dataiter)

mlflow.pytorch.log_model(
pytorch_model=model,
name=model_name, # Deprecated: artifact_path=f'/{experiment.experiment_id}/{run.info.run_id}'
#signature=signature
input_example = inputs_example.numpy()
)


Alas!, the result is an empty schema:

I also tried adding a signature instead of the input_example parameter but that didn't change anything.

I do not get any errors and the training process that follows finishes without any issues producing the model along with all other artifacts mentioned above.

Continue reading...