Solved by sunshk1227 in post #2 
              
              
                Yeah, I find a solution with the help of Szymon Maszke. Use torch.utils.data.random_split instead. Namely, 
train_data, val_data = torch.utils.data.random_split(
      train_data, (num_train, num_val))