RuntimeError: CUDA out of memory. Tried to allocate 126.00 MiB (GPU 1; 10.76 GiB total capacity; 6.98 GiB already allocated; 129.69 MiB free; 7.17 GiB reserved in total by PyTorch)

Can you try running torch.cuda.empty_cache() to free up the reserved 7.17GB memory? These reserved memory might be full of small blocks that cannot accommodate the requested 126MB.