This should work. Change the if condition in my_collate according to your needs.
def my_collate(batch):
    modified_batch = []
    for item in batch:
        image, label = item
        if label == 1:
            modified_batch.append(item)
    return default_collate(modified_batch)

train_loader = torch.utils.data.DataLoader(
    datasets.MNIST('data', train=True, download=True, transform=transform),
    batch_size=batch_size, shuffle=True, collate_fn = my_collate)