I’ve isolated the main time bottleneck of my code which seems to be in the following lines:
# subset_idx = len(train_dataset) - (some_indices)
train_dataset.train_data = train_dataset.train_data[subset_idx, :]
train_dataset.train_labels = train_dataset.train_labels[subset_idx,]

Note that the train_dataset is MNIST in my case. These 2 lines occur X numbers of times in a for loop, and they’re slowing down the loop by approximately 3.5 seconds every iteration!
I’m not using the DataLoader class, so I’m trying to see if there’s a more efficient way of reducing the total training dataset at each iteration.
Thanks!