about test ['source_cls']

Hello, thank you for your excellent work! I followed the code to complete the training process, but I encountered an issue during testing:

model_source = torch.load(config.load, map_location='cpu', weights_only=True)['source_cls'].cuda()

The saved model/checkpoint doesn’t contain the ['source_cls'] key. Could you please help clarify this? Thank you very much.