DataLoader returns the batch as a list

use your collect_fn in dataloader to return Tensor(of tensor), instead or list of tensor
look at Looping over DataLoader Returns List? - #3 by Imahn