Issue Logging SetFitModel with MLflow
I’m encountering an issue when trying to log aSetFitModelwith MLflow. MySetFitModelis loaded from a.pthfile using PyTorch’storch.load()function. However, when I try to log the model withmlflow.pytorch.log_model(), I’m receiving a TypeError stating that the ‘pytorch_model’ argument should be atorch.nn.Module.import torch
import mlflow

model_path = "/path/to/my/model.pth"
model_V5 = torch.load(model_path, map_location=torch.device('cpu'))
model_V5.to('cpu')

mlflow.start_run()
mlflow.pytorch.log_model(model_V5, "my_SetFit_model_V5")Any guidance on how to log aSetFitModelwith MLflow would be greatly appreciated.Thank you in advance!