diff --git a/hft.py b/hft.py
index 7a43bad..cc0df26 100644
--- a/hft.py
+++ b/hft.py
@@ -17,7 +17,6 @@ os.environ['TORCH_USE_CUDA_DSA'] = '1'
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 login(token="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")
 
-# Nowe tokeny specjalne
 CITATION_START = "▌▌CITATION_START"
 CITATION_END = "▌▌CITATION_END"
 
@@ -25,15 +24,15 @@ class SourceMapper:
     def __init__(self):
         self.source_to_idx = defaultdict(lambda: len(self.source_to_idx))
         self.idx_to_source = {}
-        
+
     def add_source(self, source):
         if source and source not in self.source_to_idx:
             idx = self.source_to_idx[source]
             self.idx_to_source[idx] = source
-            
+
     def get_idx(self, source):
         return self.source_to_idx[source] if source else -1
-    
+
     def get_source(self, idx):
         return self.idx_to_source.get(idx, "Unknown")
 
@@ -180,7 +179,6 @@ class CustomModel(nn.Module):
         self.base_model = AutoModelForCausalLM.from_pretrained(model_name, config=config)
         self.source_embedding = nn.Embedding(10000, config.hidden_size, padding_idx=-1)
         
-        # Dodaj specjalne tokeny i zaktualizuj embeddings
         tokenizer.add_special_tokens({'additional_special_tokens': [CITATION_START, CITATION_END]})
         self.base_model.resize_token_embeddings(len(tokenizer))
         
@@ -215,7 +213,7 @@ class CustomDataCollator(DataCollatorForLanguageModeling):
         batch = super().torch_call(examples)
         
         if "source_idx" in examples[0]:
-            source_idx = torch.tensor([ex["source_idx"] for ex in examples])
+            source_idx = torch.stack([ex["source_idx"] for ex in examples])
             batch["source_idx"] = source_idx
             
         return batch
@@ -224,11 +222,9 @@ def main():
     source_mapper = SourceMapper()
     model_name = "crumb/nano-mistral"
     
-    # Inicjalizacja tokenizera
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     tokenizer.pad_token = tokenizer.eos_token
     
-    # Przygotowanie danych
     catalog_path = "catalog.json"
     data = prepare_dataset("docs", catalog_path, source_mapper)
     
@@ -246,16 +242,18 @@ def main():
             max_length=512,
             return_tensors="pt"
         )
+        
+        source_idx = torch.tensor(examples["source_idx"], dtype=torch.long)
+        
         return {
             "input_ids": tokenized["input_ids"].squeeze(),
             "attention_mask": tokenized["attention_mask"].squeeze(),
             "labels": tokenized["input_ids"].squeeze().clone(),
-            "source_idx": examples["source_idx"]
+            "source_idx": source_idx
         }
 
     tokenized_dataset = dataset.map(tokenize_function, batched=True, batch_size=16)
 
-    # Inicjalizacja modelu z tokenizerem
     model = CustomModel(model_name, tokenizer)
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     model.to(device)