fix

src/instructlab/training/utils.py

@@ -148,14 +148,13 @@ def pad_collate_fn(batch):
 
             for idx in range(len(batch)):
                 input_ids += batch[idx]["input_ids"]
-                if has_label:
-                    labels += [-100] + batch[idx]["labels"][1:]
-                else:
-                    labels += [-100] + batch[idx]["input_ids"][1:]
+                _label = batch[idx]["labels"] if has_label else batch[idx]["input_ids"].clone()
+                _label[0] = -100
+                labels += _label
                 position_ids += list(range(len(batch[idx]["input_ids"])))
 
-            input_ids = torch.cat([x['input_ids'] for x in batch]).unsqueeze(0)
-            labels = torch.cat([x['labels'] for x in batch]).unsqueeze(0)
+            input_ids = torch.cat(input_ids).unsqueeze(0)
+            labels = torch.cat(labels).unsqueeze(0)
             position_ids = torch.tensor(position_ids, dtype=torch.long).unsqueeze(0)
 
             num_loss_counted_tokens = (labels != -100).sum()