from torch.utils.data import TensorDataset, DataLoader
BATCH_SIZE = 512
tensor_x = torch.Tensor(images[train_indices])
tensor_y = torch.Tensor(labels[train_indices])
dataset_train = TensorDataset(tensor_x, tensor_y)
train_loader = DataLoader(dataset_train, batch_size=BATCH_SIZE)
you can iterate
for i, (x, y) in enumerate(train_loader):
print(“x.shape {}, y.shape {}”.format(x.shape, y.shape))