Thanks a lot @ptrblck for your awesome help! I did not know of Subset which totally did the trick in one line. I am adding below the solution for future reference.
def _get_train_data_loader(batch_size, training_dir, is_distributed, indices, **kwargs):
    logger.info("Get train data loader")
    
    # Train = True, indices[:55000]
    dataset = datasets.MNIST(training_dir, train=True, transform=transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ]))
    
    train_ds  = torch.utils.data.Subset(dataset, indices[:55000])
    train_sampler = torch.utils.data.distributed.DistributedSampler(train_ds) if is_distributed else None
    train_loader = torch.utils.data.DataLoader(train_ds, batch_size=batch_size, shuffle=train_sampler is None, sampler=train_sampler, **kwargs)
    return train_loader

def _get_valid_data_loader(valid_batch_size, training_dir, indices, **kwargs):
    logger.info("Get valid data loader")
    
    # train = True, indices[55000:]
    dataset = datasets.MNIST(training_dir, train=True, transform=transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ]))
    
    valid_ds  = torch.utils.data.Subset(dataset, indices[55000:])
    valid_loader = torch.utils.data.DataLoader(valid_ds, batch_size=valid_batch_size, shuffle=True, **kwargs)
    return valid_loader

def _get_test_data_loader(test_batch_size, training_dir, **kwargs):
    logger.info("Get test data loader")
    
    # train = False
    dataset = datasets.MNIST(training_dir, train=False, transform=transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ]))
    
    test_loader= torch.utils.data.DataLoader(dataset, batch_size=test_batch_size, shuffle=True, **kwargs)
    return test_loader 

Note that I am calling these function from the train function like this:
    indices = np.arange(60000)
    np.random.shuffle(indices)    
    train_loader = _get_train_data_loader(args.batch_size, args.data_dir, is_distributed, indices, **kwargs)
    valid_loader = _get_valid_data_loader(args.valid_batch_size, args.data_dir, indices, **kwargs)
    test_loader = _get_test_data_loader(args.test_batch_size, args.data_dir, **kwargs)```