Call encode() with GPU

[victorcwai]

It would be great if there is an option to call model.encode("squareslab") with GPU instead of CPU.

If we call model.to(torch.device("cuda")), it doesn't work because return_dict is on CPU:

VarCLR/varclr/models/encoders.py

Line 121 in 2fe9595

self, return_dict["input_ids"], return_dict["attention_mask"]

[qibinc]

Hi @victorcwai ,

You can add

return_dict = dict(
    input_ids=return_dict["input_ids"].to("cuda")
    attention_mask=return_dict["attention_mask"].to("cuda")
)

[victorcwai]

Yes thank you.
But in my opinion, it would also be great if we are able to choose whether to use CPU and GPU in the encode() API. :-)