fix to float32

turbo-llm · Oct 21, 2024 · 85d5045 · 85d5045
1 parent 34e2895
commit 85d5045
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/turbo_alignment/dataset/chat/chat.py b/turbo_alignment/dataset/chat/chat.py
@@ -307,7 +307,7 @@ def _encode(
 
             encoded_record: dict[str, Any] = {
                 # 'id': record.id, FIXME: dont work with collators
-                'input_ids': torch.LongTensor(input_ids.astype(np.int64)),
+                'input_ids': torch.LongTensor(input_ids.astype(np.float32)),
                 'labels': torch.LongTensor(labels),
                 'attention_mask': torch.ones(input_ids.shape, dtype=torch.int64),
             }