Dataset not pulling with custom class

I tried coding a custom data set. When I try to load the class into a data loader, it states that:

NameError Traceback (most recent call last)
16 # Generators
17 training_set = roof_dataset(partition[‘train’], labels, transform = train_transforms)
—> 18 training_generator = data.DataLoader(training_set, **params)
20 test_set = roof_dataset(partition[‘test’], labels, transform = train_transforms)

NameError: name ‘data’ is not defined

Here is the code pulling the data:

CUDA for PyTorch

use_cuda = torch.cuda.is_available()
device = torch.device(“cuda:0” if use_cuda else “cpu”)


params = {‘batch_size’: 1000,
‘shuffle’: True,
‘num_workers’: 4}
max_epochs = 100


training_set = roof_dataset(partition[‘train’], labels, transform = train_transforms)
training_generator = data.DataLoader(training_set, **params)

Here is the custom class. I’m not sure where it’s going wrong;

root_dir = ‘D:\CIS inspection images 0318\train\roof\’
class roof_dataset(Dataset):

'Characterizes a dataset for PyTorch'

def __init__(self, list_IDs, labels, transform):
    self.labels = labels
    self.list_IDs = list_IDs
    self.transform = transform

def __len__(self):
    'Denotes the total number of samples'
    return len(self.list_IDs)

def __getitem__(self, index):

        'Generates one sample of data'
        # Select sample
        ID = self.list_IDs[index]

        # Load data and get label
        X = torch.load(Path(root_dir) + ID)
        y = self.labels[ID]
        if self.transform:
            X = self.transform(X)

        return X, y

From the error that you linked, it seems that you’re missing a from torch.utils import data to have data.DataLoader to work properly.