tensor_dataset.py

def get_dataloader_from_pth(path, batch_size=4):
    contents = torch.load(path)
    dataset = torch.utils.data.TensorDataset(contents['x'], contents['y'])
    dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size,
                                          shuffle=True, num_workers=2)
    return dataloader
    
#----------------------------------------------------------------------

import os.path as osp

datasetpath = '<<<ROOT>>>'
train_pth = osp.join(datasetpath, 'train.pth')
val_pth = osp.join(datasetpath, 'val.pth')
test_pth = osp.join(datasetpath, 'test.pth')

#------------------------------------------------------------------------

trainloader = get_dataloader_from_pth(train_pth)
valloader = get_dataloader_from_pth(val_pth)
testloader = get_dataloader_from_pth(test_pth)