Hi @NoobCoder,
@ptrblck is suggesting something like this,
class YourCustomSubset(torch.utils.data.Subset):
    r"""
    Same as torch.utils.data.Subset. This outputs the indices of each batch, you can choose which indices to keep.
    """

    def __getitem__(self, idx):
        return self.dataset[self.indices[idx]], idx

Check this gist/Colab notebook for the full implementation with dummy example.