I am trying to divide MNIST into separate subsets where each subset is a different class.
train_dataset = datasets.MNIST('/some/dir/', train=True, download=True, transform=self.dataset_transform)

for i in range(10):
    idx = train_dataset.targets==i
    train_dataset.targets[idx]                
    train_subsets.append(torch.utils.data.Subset(train_dataset, idx))
    train_loaders.append(torch.utils.data.DataLoader(train_subsets[i], batch_size=self.batch_size, shuffle=True))
    train_iters.append(iter(train_loaders[i]))

Beginning of my training loop:
for label_id in range(10):
    data, target = train_iters[label_id].next()

When training, I get a ValueError:
only one element tensors can be converted to Python scalars

What is going on here?