nikhil6041:

class CustomDataset(Dataset):
	def __init__(self, root_dir,  transforms=None):
	        self.root_dir = root_dir
		self.transforms = transforms
		self.img_list = sorted(glob.glob(root_dir + "/*/*"))

	def __len__(self):
		return len(self.img_list)

	def __getitem__(self, idx):
		if torch.is_tensor(idx):
		    idx = idx.tolist()	
		img_name = self.img_list[idx]
		image = io.imread(img_name)
		if self.transforms:
		    sample = self.transforms(image)	
		return image



@ptrblck
what if i change my custom dataset to return the labels in that case?