refactor(dump_sft_file.py， utils.py): 优化内存占用

2025-07-22 13:30:53 +08:00
parent 5ea1c6355d
commit 64222aeab1
2 changed files with 11 additions and 9 deletions
@@ -13,8 +13,8 @@ def get_processor(tokenizer: BpeTokenizer):
        
        tokens = prefix_ids + suffix_ids
        tokens = torch.tensor(tokens, dtype=torch.int32)
-        masks = torch.zeros_like(tokens)
-        masks[:len(prefix_ids)] = 1
+        masks = torch.zeros_like(tokens, dtype=torch.bool)
+        masks[:len(prefix_ids)] = True
        
        return {"sequence": tokens, "mask": masks}