Hi,
The problem comes from the fact that your Y_train is a 1D tensor and thus when the batch is created, its stacking plain numbers (creating a double tensor to have best precision)
Reshaping your Y_train to a 2D tensor solved the problem:
Y_train = torch.from_numpy(Y_train).view(-1, 1)

@apaszke changing this 68 line with:
 return self.data_tensor.narrow(0, index, 1), self.target_tensor.narrow(0, index, 1)

should solve the issue by always returning Tensor and not numbers. Would this break something I’m not aware of? (Let me know if you want me to send a PR for that)