Python API Transformer.from_pretrained support directly to load on GPU

❓ Questions and Help

Hi,

I'm using the API command to load my finetuned mBART model, below is the command:

en2ar = TransformerModel.from_pretrained(
  '/home/ubuntu/mayub/models',
  checkpoint_file='checkpoint_best_en-ar_ft.pt',
  data_name_or_path='/home/ubuntu/mayub/models',
  bpe='sentencepiece',
  sentencepiece_vocab='sentencepiece.bpe.model'
)

Looks like this loads the model first on cpu and then I do en2ar.cuda() to replicate again on GPU device.

Is there a way I can avoid the loading on both devices and just directly load on GPU. Reason is I have 5 other models which I need to load and save on RAM.

Appreciate any help. Thanks !