feat(trainer): 重构数据集与策略模块以支持字典形式的数据返回

2025-09-27 14:11:27 +08:00
parent 9fbc9481b5
commit 4fcdc87c95
3 changed files with 97 additions and 71 deletions
@@ -46,10 +46,10 @@ def test_env():
            return self.length
            
        def __getitem__(self, idx):
-            return (
-                torch.randint(0, 1000, (64,)),
-                torch.randint(0, 1000, (64,))
-            )
+            return {
+                "input_ids": torch.randint(0, 1000, (64,)),
+                "target_ids": torch.randint(0, 1000, (64,))
+            }
    
    dataset = DummyDataset()