The device problem when running the ESM-IF

Hi, there is the problem. How can I solve it?
ubuntu@ip:~/esm/examples/inverse_folding$ python sample_sequences.py data/5YH2.pdb     --chain C --temperature 1 --num-samples 3     --outpath output/sampled_sequences.fasta
Downloading: "https://dl.fbaipublicfiles.com/fair-esm/models/esm_if1_gvp4_t16_142M_UR50.pt" to /home/ubuntu/.cache/torch/hub/checkpoints/esm_if1_gvp4_t16_142M_UR50.pt
/home/ubuntu/.conda/envs/tcrbert/lib/python3.9/site-packages/esm/pretrained.py:215: UserWarning: Regression weights not found, predicting contacts will not produce correct results.
  warnings.warn(
Transferred model to GPU
Native sequence loaded from structure file:
SVLQSLFEHPLYRTVLPDLTEEDTLFNLNAEIRLYPKAASESYPNWLRFHIGINRYELYSRHNPVIAALLRDLLSQKISSVGMKSGGTQLKLIMSFQNYGQALFKPMKQTREQETPPDFFYFSDFERHNAEIAAFHLDRILDFRRVPPVAGRLVNMTREIRDVTRDKKLWRTFFVSPANNICFYGECSYYCSTEHALCGKPDQIEGSLAAFLPDLALAKRKTWRNPWRRSYHKRKKAEWEVDPDYCDEVKQTPPYDRGTRLLDIMDMTIFDFLMGNMDRHHYETFEKFGNDTFIIHLDNGRGFGKHSHDEMSILVPLTQCCRVKRSTYLRLQLLAKEEYKLSSLMEESLLQDRLVPVLIKPHLEALDRRLRLVLKVLSDCVEKDGFSAVVENDLD
Saving sampled sequences to output/sampled_sequences.fasta.

Sampling.. (1 of 3)
Traceback (most recent call last):
  File "/home/ubuntu/esm/examples/inverse_folding/sample_sequences.py", line 124, in <module>
    main()
  File "/home/ubuntu/esm/examples/inverse_folding/sample_sequences.py", line 120, in main
    sample_seq_singlechain(model, alphabet, args)
  File "/home/ubuntu/esm/examples/inverse_folding/sample_sequences.py", line 34, in sample_seq_singlechain
    sampled_seq = model.sample(coords, temperature=args.temperature, device=torch.device('cuda'))
  File "/home/ubuntu/.conda/envs/tcrbert/lib/python3.9/site-packages/esm/inverse_folding/gvp_transformer.py", line 123, in sample
    logits, _ = self.decoder(
  File "/home/ubuntu/.conda/envs/tcrbert/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/ubuntu/.conda/envs/tcrbert/lib/python3.9/site-packages/esm/inverse_folding/transformer_decoder.py", line 117, in forward
    x, extra = self.extract_features(
  File "/home/ubuntu/.conda/envs/tcrbert/lib/python3.9/site-packages/esm/inverse_folding/transformer_decoder.py", line 158, in extract_features
    positions = self.embed_positions(
  File "/home/ubuntu/.conda/envs/tcrbert/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/ubuntu/.conda/envs/tcrbert/lib/python3.9/site-packages/esm/modules.py", line 276, in forward
    return self.weights.index_select(0, positions.view(-1)).view(bsz, seq_len, -1).detach()
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument index in method wrapper__index_select)
Thank you very much.
Best regards
facebookresearch / esm

The device problem when running the ESM-IF #487
