diff --git a/hft.py b/hft.py
index d58f531..7ca3dcb 100644
--- a/hft.py
+++ b/hft.py
@@ -12,9 +12,11 @@ import json
 from collections import defaultdict
 from huggingface_hub import login
 
-login(token="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")
+os.environ['TORCH_USE_CUDA_DSA'] = '1'
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
+login(token="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")
+
 class SourceMapper:
     def __init__(self):
         self.source_to_idx = defaultdict(lambda: len(self.source_to_idx))
@@ -110,7 +112,7 @@ def custom_collate_fn(batch):
     labels = torch.stack([torch.tensor(b["labels"]) for b in batch])
     
     source_idx = torch.tensor([b.get("source_idx", -1) for b in batch], dtype=torch.long)
-    print("source_idx shape:", source_idx.shape)  # Debugowanie
+    #print("source_idx shape:", source_idx.shape)  # Debugowanie
     return {"input_ids": input_ids, "attention_mask": attention_mask, "labels": labels, "source_idx": source_idx}
 
 class CustomModel(nn.Module):
@@ -125,8 +127,10 @@ class CustomModel(nn.Module):
         
     def forward(self, input_ids=None, attention_mask=None, labels=None, source_idx=None, **kwargs):
         if source_idx is not None:
+            print("Max source_idx:", torch.max(source_idx))
+            print("Num embeddings:", self.source_embedding.num_embeddings)
+            source_idx = torch.clamp(source_idx, 0, self.source_embedding.num_embeddings - 1)
             source_embeds = self.source_embedding(source_idx).unsqueeze(1).expand(-1, input_ids.size(1), -1)
-            # Dodaj embeddingi źródła do wejścia modelu
             hidden_states = self.base_model.get_input_embeddings()(input_ids) + source_embeds
             outputs = self.base_model(inputs_embeds=hidden_states, attention_mask=attention_mask, labels=labels, **kwargs)
         else:
@@ -158,7 +162,7 @@ tokenized_dataset = dataset.map(tokenize_function, batched=True, batch_size=8)
 config = AutoModelForCausalLM.from_pretrained(model_name).config
 print("Vocabulary size:", config.vocab_size)
 model = CustomModel(model_name, config)
-model.to("cpu")
+model.to("cpu")  # Zmienione na CPU dla debugowania
 
 # Konfiguracja treningu
 training_args = TrainingArguments(
@@ -167,7 +171,7 @@ training_args = TrainingArguments(
     per_device_train_batch_size=2,
     gradient_accumulation_steps=4,
     learning_rate=2e-5,
-    fp16=True,
+    fp16=False,  # Wyłączone dla CPU
     logging_steps=1,
     logging_dir="./logs",
     save_strategy="steps",
@@ -199,12 +203,4 @@ def generate_answer(question, model, tokenizer, source_mapper, max_length=200):
     
     # Pobierz źródło z ostatniego tokena
     last_token_id = outputs.sequences[0][-1].item()
-    source_idx = model.source_embedding.weight.shape[0] - 1  # Tymczasowe rozwiązanie
-    source = source_mapper.get_source(source_idx)
-    
-    return f"{answer}\n\nŹródło: {source if source else 'Opracowanie własne'}"
-
-# Przykład użycia
-question = "Ile dni urlopu przysługuje pracownikowi?"
-answer = generate_answer(question, model, tokenizer, source_mapper)
-print(answer)
+    source_idx = model.source_embeddi