Hi PyTorch community, I hope this is the right section to post my issue in.
I have a quite big dataset of biological (protein) sequences, around 68M, that I have to process with Rostlab/prot_t5_xl_uniref50 at main 2, a pre-trained NLP model that generates embeddings for downstream tasks. I have set up a decent workflow that is able to recover when my single V100 GPU goes OOM when the sequence is too big (even with batch size set to 1), effectively resorting to CPU and RAM for the calculations.
As of now, I used the suggested PyTorch implementation in the for loop:
oom = False
try:
    embedding = model(input_ids=input_ids, attention_mask=attention_mask)
except RuntimeError:
    oom = True
if oom:
    model.cpu()
    embedding = model(input_ids=input_ids.cpu(), attention_mask=attention_mask.cpu())

This approach solves the OOM issues, but I realized that has the side effect of having the GPU sitting around doing nothing while the big sequence is being processed, only to resume working with the next iterations.
Is there any possibility to send the problematic sequence to CPU while continuing to process the rest normally?
Things I’ve tried:

DataParallel is not the right answer:ValueError: Expected a non cpu device, but got: cpu

Python multiprocessing expects a pickleable Dataset: AttributeError: Can't pickle local object 'main.<locals>.SequenceDataset'