@ptrblck the data loader is
def data_loaders():
    train_data = datasets.MNIST(
        root = 'data',
        train = True,                         
        transform = transforms.ToTensor(), 
        download = True,            
    )
    test_data = datasets.MNIST(
        root = 'data', 
        train = False, 
        transform = transforms.ToTensor()
    )

    train_test_loaders = {
        'train' : torch.utils.data.DataLoader(train_data, 
                                            batch_size=100, 
                                            shuffle=True, 
                                            num_workers=1),
        
        'test'  : torch.utils.data.DataLoader(test_data, 
                                            batch_size=100, 
                                            shuffle=True, 
                                            num_workers=1),
    }
    return train_test_loaders

Later, at the time of training
def train(NB_EPOCS, model, loaders):
    model.train()
        
    # Train the model
    total_step = len(loaders['train'])
    for epoch in range(NB_EPOCS):
        for i, (images, labels) in enumerate(loaders['train']):
            b_x = Variable(images)   # batch x
            b_y = Variable(labels)   # batch y
            print(np.shape(b_x), np.shape(b_y))
            output = model(b_x)[0]               
            loss = criterion(output, b_y)

The shape of the b_x and b_y is
torch.Size([100, 1, 28, 28]) torch.Size([100])