add_to_index uses too much GPU RAM and crashes

Hi there,

I am implementing Ragatouille at scale for a RAG model project. I am running the add_to_index() method sequentially but I keep running into memory issues with my GPU with 16 GB RAM. Does the add_to_index() method use GPU RAM or am I implementing something incorrectly?

Is there a way to optimize my GPU usage? Can I also run only on CPU?

Thanks.

Martin

Provided code:

from ragatouille import RAGPretrainedModel
RAG = RAGPretrainedModel.from_pretrained("colbert-ir/colbertv2.0")

RAG.index(
collection=[file_content],
document_metadatas=[{'source': 'initial_data.txt', 'page': 0}],
document_ids=[str(uuid.uuid4())],
index_name="test3",
split_documents=True,
max_document_length=300,
#use_faiss=True
)

for _ in range(10):
RAG.add_to_index(
new_collection=all_splits_rgt,
new_document_metadatas=rgt_metadatas,
new_document_ids=[str(uuid.uuid4()) for _ in range(len(all_splits_rgt))],
#index_name="test1",
#new_split_documents=True,
#new_max_document_length=300,
#use_faiss=True
)
retriever_advanced_text = RAG_redone.as_langchain_retriever(k=10)