How to reduce memory usage of ESM2-650

[YuAoZeng]

My code is as follows：

output:
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 7.53 GiB. GPU 0 has a total capacty of 47.54 GiB of which 6.75 GiB is free. Including non-PyTorch memory, this process has 40.75 GiB memory in use. Of the allocated memory 40.32 GiB is allocated by PyTorch, and 121.56 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting max_split_size_mb to avoid fragmentation. See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

I want to use ESM-2 to get the representation of protein sequences, but I only input 4 protein sequences with a length of no more than 1000, and the memory is not enough. Is this normal? How can I reduce memory usage?

[yamule]

Did you try return_contacts=False ?

[YuAoZeng]

I haven't tried it, can you teach me how to use return_contacts=False？

[yamule]

Just change this line
results = model(batch_tokens, repr_layers=[33], return_contacts=True)
as follows.
results = model(batch_tokens, repr_layers=[33], return_contacts=False)

& check that the results are as you expected.

[YuAoZeng]

Thank you very much for your help, I solved it