Merge pull request !2363 from yanghaitao/yht_fix_textfiledatasettags/v0.5.0-beta
| @@ -1001,7 +1001,7 @@ class Dataset: | |||||
| if isinstance(sampler, samplers.DistributedSampler): | if isinstance(sampler, samplers.DistributedSampler): | ||||
| dev_id = sampler.shard_id | dev_id = sampler.shard_id | ||||
| return "", dev_id | return "", dev_id | ||||
| if isinstance(output_dataset, TFRecordDataset): | |||||
| if isinstance(output_dataset, (TFRecordDataset, TextFileDataset, CLUEDataset)): | |||||
| if output_dataset.shard_id is not None: | if output_dataset.shard_id is not None: | ||||
| dev_id = output_dataset.shard_id | dev_id = output_dataset.shard_id | ||||
| return "", dev_id | return "", dev_id | ||||
| @@ -344,6 +344,15 @@ def test_clue_wsc(): | |||||
| }) | }) | ||||
| assert len(buffer) == 3 | assert len(buffer) == 3 | ||||
| def test_clue_to_device(): | |||||
| """ | |||||
| Test CLUE with to_device | |||||
| """ | |||||
| TRAIN_FILE = '../data/dataset/testCLUE/afqmc/train.json' | |||||
| data = ds.CLUEDataset(TRAIN_FILE, task='AFQMC', usage='train', shuffle=False) | |||||
| data = data.to_device() | |||||
| data.send() | |||||
| if __name__ == "__main__": | if __name__ == "__main__": | ||||
| test_clue() | test_clue() | ||||
| @@ -89,6 +89,10 @@ def test_textline_dataset_get_datasetsize(): | |||||
| size = data.get_dataset_size() | size = data.get_dataset_size() | ||||
| assert size == 3 | assert size == 3 | ||||
| def test_textline_dataset_to_device(): | |||||
| data = ds.TextFileDataset(DATA_FILE, shuffle=False) | |||||
| data = data.to_device() | |||||
| data.send() | |||||
| if __name__ == "__main__": | if __name__ == "__main__": | ||||
| test_textline_dataset_one_file() | test_textline_dataset_one_file() | ||||