This PR fix the output format of the create_maestro_datasets function.
1. Motivation and Context
Originally, the dattri/benchmark/maestro/data/create_maestro_datasets function will directly output three dataloaders. But I found that the current dattri/scripts/dattri_retrain.py actually will assume this function to output torch dataset objects (will construct data loaders in that script). Thus make this change.
Description
This PR fix the output format of the
create_maestro_datasets
function.1. Motivation and Context
Originally, the
dattri/benchmark/maestro/data/create_maestro_datasets
function will directly output three dataloaders. But I found that the currentdattri/scripts/dattri_retrain.py
actually will assume this function to output torch dataset objects (will construct data loaders in that script). Thus make this change.2. Summary of the change