Loading segmentation model - doesn't output masks

I am training a segmentation model and then loading it for inference.
Training (SemTorch for model):

learn = get_segmentation_learner(dls=dls, number_classes=2, segmentation_type="Semantic Segmentation",
                                 architecture_name="unet", backbone_name="resnet34", 
                                 metrics=[Dice(), JaccardCoeff()], opt_func=ranger)
learn.fit_flat_cos(20, 1e-3, cbs=SaveModelCallback(monitor='dice'), pct_start=0.5)

torch.save(dls, 'dls.pkl')

When I use show_results() after training, it shows the true and predicted masks.

Loading (in a new notebook):

learn = load_learner('learn.pkl')

When I use show_results() after loading, the model does not predict any masks - they are all blank!

However, if I train the model for one epoch, it then predicts masks (and does it well). However, I DO NOT want to train for an epoch before conducting inference, because each epoch takes more than 2 minutes on a GPU (Tesla T4).

Also I need to use SemTorch, because I am using three different architectures (UNet, HRNet and DeepLabV3+) and creating an ensemble from the three - vanilla FastAi only has a UNet.

Edit: SemTorch is not the cause. Fastai models also have this problem. Just checked.

Fixed it by using learner.save and learner.load instead of learner.export and load_learner()