Moving trained model between devices

Hello! Is there a way to move the trained SetFit model to another device? E.g. I train a model on a GPU, and after initializing the model on a CPU-only machine I get RuntimeError: Attempting to deserialize object on a CUDA device but torch.cuda.is_available() is False. If you are running on a CPU-only machine, please use torch.load with map_location=torch.device('cpu') to map your storages to the CPU.

I have managed to fix the error by running:

model.model_body = model.model_body.to("cpu")
model.model_head = model.model_head.to("cpu")

However, I guess it depends on the type of the head (in case of the Logistic regression head there is no need to move it).

If there is no single method to move a model between devices, shall I open a PR?

[JackMurphy97]

I have also come across this exact issue - it would be awesome to have some .to(device) functionality for the SetFitModel class!

[blakechi]

Thank you for your detailed explanation! I guess the error raised because of the initialization method of SetFitHead. Will validate my guess shortly.

Thanks again for raising this up!