How to load to from a checkpoint to same device when pretrained encoder was used
❓ Questions and Help
I implemented a ClassificationNet (see below) that's using a pretrained encoder. After training, I'm trying to load it to CPU using ClassificationNet.load_from_checkpoint(pth, map_location=torch.device("cpu"), but since map_location in get_encoder is None, the encoder tries to load to GPU. How can I inform get_encoder to load to the same map_location?
Since I just started using Lightning, I guess there's a much smarter way to circumvent this situation altogether -- I look forward to your suggestions :) Thanks!
Code
class ClassificationNet(LightningModule):
    ...
    self.encoder = get_encoder(pretrained=True)

get_encoder(pretrained=False, map_location=None):
    model = FancyModel()
    if pretrained:
        ckpt_data = torch.utils.model_zoo.load_url(url, map_location=map_location)
    ....

OS: Manjaro Linux
Version 1.0.5