Hi @ptrblck
I have tried to repeat your exercise which provide to me a very good dbg idea (save input/output of some reference samples and at the same time the model dictionary). To do so I have rewritten fresh codes. During the training session, after the training optimisation, I send to the model after model.eval() 10 references inputs. Then, I save the model parameters, the refernce inputs and the corresponding outputs:
        test_ref_out = test_ref(args, model, device, sample_test_ref,
                        test_transforms, epoch)

        #save model state and reference test in/out
        state = {
            'model_state_dict': model.state_dict(),
            **test_ref_out
            }

with  the returned dictionary test_ref_outdefined as
{'ref_in': new_img_batch, 'ref_ebv': ebv_batch, 'ref_out': output}

Now, in the debugging session
np.random.seed(0)
torch.manual_seed(0)
torch.cuda.manual_seed_all(0)
random.seed(0)

use_cuda = torch.cuda.is_available()

device = torch.device("cuda" if use_cuda else "cpu")
print("Use device....: ",device)

# Pin_memory speed up the CPU->GPU transfert for NVidia GPU
kwargs = {'num_workers': 0, 'pin_memory': True} if use_cuda else {}


#model
model = resnet20()
# in place method to transfet to GPU
# according to https://pytorch.org/docs/stable/nn.html#torch.nn.Module.to
model.to(device)

#load model parameters
root_file = "./"
checkpoint_file ="dbg-resnet20-2.pth"

checkpoint_file = root_file + checkpoint_file
checkpoint = torch.load(checkpoint_file)
model.load_state_dict(checkpoint['model_state_dict'])

#x-check model inference form reference samples
ref_input_img = checkpoint['ref_in']
ref_input_ebv = checkpoint['ref_ebv']
ref_output    = checkpoint['ref_out']
model.eval()
with torch.no_grad():
    output = model(ref_input_img,ref_input_ebv)
# Compare
diff = (output - ref_output).abs().max().item()
print('max abs error: ', diff)
if diff>0:
    print('ref_out\n',ref_output,'\n',F.softmax(ref_output,dim=1))
    print('new_out\n',output,'\n',F.softmax(output,dim=1))

#x-check if load the same data
test_ref_file  = 'test10k.npz'
# Reference test
test_ref_dataset = DatasetPz(test_ref_file, transform=None)
test_ref_loader = torch.utils.data.DataLoader(
    test_ref_dataset,
    batch_size=10, drop_last=True,
    shuffle=False, **kwargs)

print("test_ref_loader length=",len(test_ref_loader))
testref_loader_iterator = iter(test_ref_loader)
sample_test_ref = next(testref_loader_iterator) # once for all
img_ref = sample_test_ref['image'].transpose(3,1).transpose(3,2)  # input image

ebv_ref = sample_test_ref['ebv']    # input redenning
img_ref,ebv_ref = img_ref.to(device),ebv_ref.to(device)
# Compare
print('img: max abs error: ', (img_ref - ref_input_img).abs().max())
print('ebv: max abs error: ', (ebv_ref - ref_input_ebv).abs().max())

The result is not satisfactory for the output of the model after reloaded as it finds different values from the loaded reference inputs, while I have x-checked that the reference inputs corresponds to the same data that I can directly laod from file (2nd & 3rd comparisons)
Model output:
max abs error:  11.42324447631836

The reference values saved in the checkpoint file: (I also produce the F.softmax(out,dim=1) to get the probability distribution
ref_out
 tensor([[-1.2732, -0.4523,  0.3474,  ..., -2.2444, -1.8375, -1.7659],
        [11.1764, 12.5662, 11.4143,  ..., -2.0558, -0.6456, -2.7947],
        [ 3.6198,  4.7427,  4.8422,  ..., -1.8688, -1.0088, -1.7117],
        ...,
        [-1.1483, -1.3629, -1.4520,  ..., -2.2878, -2.3314, -2.1438],
        [ 1.8598,  2.7087,  2.9787,  ..., -1.2867, -0.7799, -0.9943],
        [-0.6644,  0.2698,  0.9862,  ..., -2.5374, -1.9054, -1.7161]],
       device='cuda:0') 
 tensor([[1.3692e-06, 3.1115e-06, 6.9227e-06,  ..., 5.1839e-07, 7.7875e-07,
         8.3651e-07],
        [2.8840e-02, 1.1577e-01, 3.6586e-02,  ..., 5.1681e-08, 2.1172e-07,
         2.4686e-08],
        [2.4761e-03, 7.6110e-03, 8.4070e-03,  ..., 1.0235e-05, 2.4187e-05,
         1.1977e-05],
        ...,
        [4.1218e-05, 3.3256e-05, 3.0422e-05,  ..., 1.3188e-05, 1.2626e-05,
         1.5231e-05],
        [1.9157e-03, 4.4774e-03, 5.8652e-03,  ..., 8.2377e-05, 1.3674e-04,
         1.1036e-04],
        [8.0213e-07, 2.0415e-06, 4.1790e-06,  ..., 1.2326e-07, 2.3189e-07,
         2.8021e-07]], device='cuda:0')

And now the model reloaded output
new_out
 tensor([[ 0.1373,  1.2026,  1.6492,  ..., -0.6598,  0.0265, -1.3969],
        [ 0.0962,  1.1430,  1.6046,  ..., -0.6021,  0.0197, -1.3420],
        [ 0.1854,  1.1745,  1.7196,  ..., -0.5364,  0.0581, -1.4119],
        ...,
        [ 0.2378,  1.2321,  1.7412,  ..., -0.5621,  0.0525, -1.4340],
        [ 0.1403,  1.1963,  1.6627,  ..., -0.6335,  0.0163, -1.3910],
        [ 0.0681,  1.1639,  1.6145,  ..., -0.7172,  0.0148, -1.3750]],
       device='cuda:0') 
 tensor([[0.0021, 0.0060, 0.0093,  ..., 0.0009, 0.0018, 0.0004],
        [0.0021, 0.0059, 0.0093,  ..., 0.0010, 0.0019, 0.0005],
        [0.0023, 0.0062, 0.0107,  ..., 0.0011, 0.0020, 0.0005],
        ...,
        [0.0024, 0.0066, 0.0109,  ..., 0.0011, 0.0020, 0.0005],
        [0.0020, 0.0059, 0.0094,  ..., 0.0009, 0.0018, 0.0004],
        [0.0018, 0.0054, 0.0084,  ..., 0.0008, 0.0017, 0.0004]],
       device='cuda:0')

So, why my model reloading does not get the same outputs as during training phase even if in both case I have used model.eval() ??? what about the random ganerator ?
I do not know if you have noticed that I systematically use at the beginning of my scripts
np.random.seed(0)
torch.manual_seed(0)
torch.cuda.manual_seed_all(0)
random.seed(0)