github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-08 12:30:42 +00:00
Code Issues Packages Projects Releases Wiki Activity

[booster] update prepare dataloader method for plugin (#3706)

Browse Source 
* [booster] add prepare dataloader method for plug

* [booster] update examples and docstr
This commit is contained in:
Hongxin Liu
2023-05-08 15:44:03 +08:00
committed by GitHub
parent f83ea813f5
commit 3bf09efe74
9 changed files with 41 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
tests/test_booster/test_plugin/test_dp_plugin_base.py
View File

@@ -55,7 +55,7 @@ def check_dataloader_sharding():
# create a custom dasetset with 0 to 10
dataset = TensorDataset(torch.arange(0, 10))
train_dataloader = plugin.prepare_train_dataloader(dataset, batch_size=2)
train_dataloader = plugin.prepare_dataloader(dataset, batch_size=2)
# get the first batch of data
batch = next(iter(train_dataloader))[0].cuda()