Summary: support classification and refactor build_dataloader which can support explicit parameters passing