[Python] TypeError: Dataloader object is not subscriptable

I'm creating an AI model to generate density plots of crowds. When splitting the dataset into two, one for training and one for validation, I create the two data sets and try to load the datasets using torch.utils.data.DataLoader(test_set, batch_size=batch_size, shuffle=False). After that, to test the data I iterate through and use the next function to get the next element of the dataset and then I get the TypeError. The dataset I'm using his the ShanghaiTech Crowd Counting data set from Kaggle: https://www.kaggle.com/datasets/tthien/shanghaitech

Here is the full code:

batch_size = 8
device = 'cuda:0' if torch.cuda.is_available() else 'cpu'

train_root_dir = "data/part_A/train_data/"
init_training_set = DataLoader(train_root_dir, shuffle=True)

# split part of the training set as validation set
train_size = int(0.9 * len(init_training_set))
val_size = len(init_training_set) - train_size

train_indices = list(range(train_size))
val_indices = list(range(train_size, len(init_training_set)))
train_dataset = torch.utils.data.dataset.Subset(init_training_set, train_indices)
val_dataset = torch.utils.data.dataset.Subset(init_training_set, val_indices)

train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=batch_size, shuffle=False)

test_root_dir = "data/part_A/test_data/"
test_set = DataLoader(test_root_dir, shuffle=False)
test_loader = torch.utils.data.DataLoader(test_set, batch_size=batch_size, shuffle=False)

dataiter = iter(train_loader)
ex_images, ex_dmaps, ex_n_people = next(dataiter)


# Show images and density map
plot_corresponding_pairs(ex_images, ex_dmaps)


The specific error is:

Traceback (most recent call last):
line 61, in <module>
for ex_images, ex_dmaps, ex_n_people in train_loader
TypeError: 'DataLoader' object is not subscriptable

Continue reading...