@ptrblck  Using your suggestions I was trying to create a reduced dataset for a problem at hand, however I have been encountering Cuda runtime error:device -side assert triggered repeatedly for the following code where I first use the Imagefolder to get my dataset and then reduce it on the basis of the threshold kept as 3 followed by a stratified split of the paths and labels to have my train,val and test sets which I finally use for my model. Do you see any error in this snippet ? Also i think it can be optimized but not really sure of how should i do it . Can you provide suggestions for that?

data_dir = original_cropped_dir_path

threshold = 3

transformations  = T.Compose([

    T.Resize(size = 256),

    T.ToTensor(),

    fixed_image_standardization

])

dataset = datasets.ImageFolder(data_dir, transform=transformations)

targets = torch.tensor(dataset.targets)

unique_classes, counts = torch.unique(targets, return_counts=True)

# print(unique_classes, counts)

classes_to_keep = unique_classes[counts>=threshold]

# print(len(classes_to_keep))

indices = []

for class_to_keep in classes_to_keep:

    indices.extend((targets == class_to_keep).nonzero())

# indices = torch.tensor(indices)

# print(indices)

for i in range(len(indices)):

    indices[i] = indices[i].item()

print(indices)

print(len(indices))

re_targets = targets[indices].tolist()

print(re_targets[0])

print(len(re_targets))

re_img_paths = []

for i in indices:

    re_img_paths.append(dataset.imgs[i][0])

print(re_img_paths[0])

print(len(re_img_paths))

from sklearn.model_selection import train_test_split

t_paths , test_paths , t_labels , test_labels = train_test_split(re_img_paths,re_targets,stratify = re_targets,test_size = 0.2)

train_paths , val_paths , train_labels , val_labels = train_test_split(t_paths,t_labels,stratify = t_labels,test_size = 0.2)

# print(np.unique(t_labels , return_counts=True))

# print(np.unique(val_labels , return_counts=True))

# print(np.unique(test_labels,return_counts=True))

# print(np.unique(np.array(re_targets),return_counts=True))

# sns.countplot(data = re_targets)

# sns.countplot(data = re_targets.numpy())

# from sklearn.model_selection import train_test_split

batch_size = 32

validation_size = .2

test_size = 0.1

shuffle_dataset = True

num_workers = 0 if os.name == 'nt' else 2

# if shuffle_dataset :

#     np.random.seed(seed_val)

#     np.random.shuffle(train_val_indices)

# train_indices, val_indices = train_val_indices[val_split:], train_val_indices[:val_split]

# # Creating PT data samplers and loaders:

# train_sampler = SubsetRandomSampler(train_indices)

# valid_sampler = SubsetRandomSampler(val_indices)

# test_sampler = SubsetRandomSampler(test_indices)

class LFWDataset(Dataset):

    def __init__(self,img_paths,labels,transforms):

        super(LFWDataset,self).__init__()

        self.img_paths = img_paths

        self.labels = labels

        self.transforms = transforms 

    def __getitem__(self,idx):

        path = self.img_paths[idx]

        img = Image.open(path).convert('RGB')

        img = self.transforms(img)

        label = torch.tensor(self.labels[idx])

        return img,label

    def __len__(self):

        return len(self.labels)

train_dataset = LFWDataset(train_paths,train_labels,transformations)

val_dataset = LFWDataset(val_paths,val_labels,transformations)

test_dataset = LFWDataset(test_paths,test_labels,transformations)

train_loader = torch.utils.data.DataLoader(

                                train_dataset, batch_size=batch_size, 

                                drop_last = True,

                                num_workers = num_workers

                )

validation_loader = torch.utils.data.DataLoader(

                                val_dataset, batch_size=batch_size,

                                drop_last = True,

                                num_workers = num_workers

                    )

test_loader = torch.utils.data.DataLoader(

                                test_dataset, batch_size=batch_size,

                                drop_last = True,

                                num_workers = num_workers

                    )