patrick126:

for output in outputs_max:
                predictions[sample_counter] = output
                sample_counter += 1



You are assigning the label of each sample by looping over each prediction in each batch. This is not optimal at all and although you are using batches but at each iteration, you are using predictions to have all predicted labels. But still it should be fine and I cannot see any issue regarding size mismatch, maybe a mistake in passing arguments to methods?
This will do the same job:
dataloader_valid = torch.utils.data.DataLoader(dataset_valid, 
                                               batch_size=len(dataset_valid),
                                               num_workers=0,
                                               pin_memory=True)

def test(dataset, dataloader):

    # switch to test mode
    net.eval()  

    # do not accumulate gradients (faster)
    with torch.no_grad():

        # test all batches
        for batch in dataloader:

            # get data from dataloader [ignore labels/targets as they are not used in test mode]
            inputs = batch[0]

            # move data to device
            inputs = inputs.to(device, non_blocking=True)

            # forward pass
            outputs = net(inputs)

            # store predictions
            predictions = torch.argmax(outputs, dim=1)

    return predictions