Pytorch Resnet CNN only works when test data contains all classes

Question

Hoping to get a hand with a weird CNN training issue.

I am training a Resnet classifier to predict 4 classes of image from a ~10k image dataset. The code is pretty si