You could perform the splitting beforehand.
Based on your code, it seems that using a torch.utils.data.Subset would be easier, as you would only have to provide the corresponding indices.



 Riko:

Now i have no idea how to use this in torch.utils.data


You can just iterate the DataLoader:
for data, target in train_loader:
    data = data.to(device)
    target =  target.to(device)
    optimizer.zero_grad()
    output = model(data)
    ...

It will automatically create the batches, shuffle your Dataset, and use multiprocessing to load the data (if num_workers > 0).