Inference AUROC on valid set using ModelCheckpoint saved ckpt does not equal valid AUROC from training
🐛 Bug
During training, I have a ModelCheckpoint callback that saves the top 5 models based on valid_AUC computed at the end of validation phase using multiclass_auroc() function. The callback saves the .ckpt file as epoch=X_valid_AUC=0.XXXX.ckpt. When I load the ckpt and run trainer.test() on the same validation set used in training, it returns a different value (see screenshot below). I decided to test if it would work using 1 GPU in ddp, and it works in that case. When using 4 GPUs it does not work.
My validation set size is 543 examples, but in my validation_epoch_end function I use self.log("valid_preds_size", preds.shape[0]) and it outputs 136 which is approx 543/4. Shouldn't it output 543 because validation_epoch_end has gathered all the batch dicts from every GPU? I create my preds tensor like this:
preds = torch.cat( [step[head.name]["preds"] for step in valid_outputs if head.name in step] )


CUDA:

GPU:

Tesla V100-SXM2-16GB
Tesla V100-SXM2-16GB
Tesla V100-SXM2-16GB
Tesla V100-SXM2-16GB


available:         True
version:           10.1


Packages:

numpy:             1.18.1
pyTorch_debug:     False
pyTorch_version:   1.7.1+cu101
pytorch-lightning: 1.1.1
tqdm:              4.46.1


System:

OS:                Linux
architecture:

64bit



processor:
python:            3.7.6
version:           #1 SMP Debian 4.9.210-1+deb9u1 (2020-06-07)



Additional context