Train, test, validation streamlined

I have successfully managed to break down the dataset into three parts ( train, test and oval). What is bugging me is how long was the way I found to do it using data loader. No need to say, I am a beginner user of PyTorch. What i am wondering is if it does not exist a more straightforward method like Sklearn, train test split.

Please see my code below.

#load the data from the path the data is only separated into three classes
train_dataset = torchvision.datasets.ImageFolder(
root= data_path,
transform = data_transforms,
target_transform=None)

# Load the data as one large batch
loader = DataLoader(
train_dataset,
batch_size=1000,
num_workers=0,
shuffle=True)

# Separate into train val
trainloader, valloader = torch.utils.data.random_split(loader.dataset, (144 , 48))

#Load train and val
 train_loader = DataLoader( #has 144
trainloader,
batch_size=16,
num_workers=0,
shuffle=True)

val_loader = DataLoader( # has 48
valloader,
batch_size=1,
num_workers=0,
shuffle=False)

#Separate val into test and val
testloader, valloader = torch.utils.data.random_split(val_loader.dataset, (40, 8))

#Load test and val
test_loader = DataLoader( # has 40
testloader,
batch_size=1,
num_workers=0,
shuffle=False)

val_loader = DataLoader( # has 8
valloader,
batch_size=1,
num_workers=0,
shuffle=False)

You could use sklearn.model_selection.train_test_split directly by using the dataset indices (just create the indices using np.arange).
Once you have the indices for each dataset split, you could either use Subset and wrap the original dataset with the corresponding indices, or use a SubsetRandomSampler and pass these samplers to each DataLoader.