CUDA out of memory when extract.py was run

[HelloYiHan]

Hi all,

When I ran the code below:

python extract.py esm2_t48_15B_UR50D example.fasta /example_out/ --include mean per_tok --toks_per_batch 2

I got the error message:

RuntimeError: CUDA out of memory.

I was wondering how to compute embeddings of bulk FASTA sequences with esm2_t48_15B_UR50D model.

Thank you for your help!

Yi

[tomsercu]

We provide an example to get you started with cpu-offloading, but it is definitely more advanced: https://github.com/facebookresearch/esm#fsdp

Let us know if you get stuck!