I cannot reproduce the issue using your code snippet and models.resnet18() as the pretrained model after removing undefined code parts. Could you update your code and make it executable to that we could reproduce the error?