I am trying to fine tune pretrained Rag on my Custom dataset. I am enclosing the code. During Training, I am getting OutOfMemoryError: CUDA out of memory. Tried to allocate 20.00 MiB. GPU 0 has a total capacity of 22.19 GiB of which 3.50 MiB is free. Process 15094 has 22.18 GiB memory in use. Of the allocated memory 21.70 GiB is allocated by PyTorch, and 195.58 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation. See documentation for Memory Management (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)
Output is truncated. View as a scrollable element or open in a text editor. Adjust cell output settings… I tried everything.
I am enclosing my code also any help would be appreciated.

# %%
!pip install transformers datasets torch

# %%
from transformers import RagTokenizer, RagTokenForGeneration
tokenizer = RagTokenizer.from_pretrained("facebook/rag-token-nq")
model = RagTokenForGeneration.from_pretrained("facebook/rag-token-nq")

# %%
!pip install accelerate -U

# %%
!pip install faiss-cpu

# %%
!pip install accelerate -U

# %%
pip show accelerate

# %%
import os

# Set the environment variable
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"


# %%
import pandas as pd
from sklearn.model_selection import train_test_split
from datasets import Dataset, DatasetDict, load_from_disk
from transformers import Trainer, TrainingArguments, RagTokenizer, RagSequenceForGeneration, RagRetriever, DPRContextEncoder, DPRContextEncoderTokenizerFast
import faiss
import torch
import numpy as np


# Load your CSV data
df = pd.read_csv("/teamspace/studios/this_studio/examples/qa_pairs1.csv")  # Replace with the path to your CSV file

# Split the data into train and validation datasets
train_df, val_df = train_test_split(df, test_size=0.2)

# Convert DataFrame to Hugging Face Dataset
train_dataset = Dataset.from_pandas(train_df)
val_dataset = Dataset.from_pandas(val_df)

# Combine into a DatasetDict
climate_change_dataset = DatasetDict({
    "train": train_dataset,
    "validation": val_dataset
})

# Generate embeddings for your passages
ctx_encoder = DPRContextEncoder.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base")
ctx_tokenizer = DPRContextEncoderTokenizerFast.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base")

def create_embeddings(sentences):
    embeddings = []
    for passage in sentences:
        inputs = ctx_tokenizer(passage, return_tensors="pt", padding=True, truncation=True, max_length=512)
        with torch.no_grad():
            embedding = ctx_encoder(**inputs).pooler_output.cpu().numpy()
        embeddings.append(embedding)
    return np.concatenate(embeddings, axis=0).astype("float32")

# Prepare the passages data
passages = df["sentence"].tolist()
embeddings = create_embeddings(passages)

# Structure your data
data = {
    "title": [""] * len(df),  # Adding dummy title column
    "text": df["sentence"].tolist(),
    "embeddings": embeddings.tolist()
}

# Create a dataset using the Dataset class from datasets
passages_dataset = Dataset.from_dict(data)

# Save the dataset to disk
passages_dataset.save_to_disk("/teamspace/studios/this_studio/examples/my_knowledge_dataset")

# Create and save FAISS index
index = faiss.IndexFlatL2(embeddings.shape[1])
index.add(embeddings)
faiss.write_index(index, "/teamspace/studios/this_studio/examples/faiss_index")

# Load tokenizers and models
question_encoder_tokenizer = RagTokenizer.from_pretrained("facebook/rag-token-nq")
generator_tokenizer = RagTokenizer.from_pretrained("facebook/rag-token-nq")

# Define the tokenizer for question encoder and generator
tokenizer = question_encoder_tokenizer

# Define a function to process the data for training
def preprocess_function(examples):
    inputs = [q + " [SEP] " + s for q, s in zip(examples["question"], examples["sentence"])]
    model_inputs = tokenizer(inputs, max_length=512, truncation=True, padding="max_length")
    labels = tokenizer(examples["answer"], max_length=512, truncation=True, padding="max_length")
    model_inputs["labels"] = labels["input_ids"]
    return model_inputs

# Process the dataset
tokenized_datasets = climate_change_dataset.map(preprocess_function, batched=True)

# Initialize the Trainer
training_args = TrainingArguments(
    output_dir="./rag_finetuned/",
    num_train_epochs=3,
    per_device_train_batch_size=4,
    per_device_eval_batch_size=4,
    warmup_steps=500,
    weight_decay=0.01,
    logging_dir="./logs/",
    logging_steps=10,
)

# Load the dataset from disk
dataset_path = "/teamspace/studios/this_studio/examples/my_knowledge_dataset"
dataset = load_from_disk(dataset_path)

# Load the FAISS index
index_path = "/teamspace/studios/this_studio/examples/faiss_index"
index = faiss.read_index(index_path)

retriever = RagRetriever.from_pretrained(
    "facebook/rag-token-nq",
    index_name="custom",
    passages_path=dataset_path,
    index_path=index_path,
    question_encoder_tokenizer=question_encoder_tokenizer,
    generator_tokenizer=generator_tokenizer,
)

model = RagSequenceForGeneration.from_pretrained(
    "facebook/rag-token-nq",
    retriever=retriever,
    question_encoder_tokenizer=question_encoder_tokenizer,
    generator_tokenizer=generator_tokenizer,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_datasets["train"],
    eval_dataset=tokenized_datasets["validation"],
    
)

# Start fine-tuning
trainer.train()


# %%
torch.cuda.empty_cache()

I am getting to fine tuene my Rag Model but getting out of Memory inspite of using GPU.