How to remove FeatureLoss from exported superres model

In lesson 7, in the superres.ipynb, we define a FeatureLoss Class that we use to train our model. Upon exporting this model with learn.export, and trying to load it with load_learner in a new notebook, I get the following error :

AttributeError: Can't get attribute 'FeatureLoss' on <module '__main__'>

My understanding of this is that the FeatureLoss of the learner is saved as well, and when trying to load it, it raises this error because FeatureLoss Class is not defined in my new notebook.
One way around this is to import the FeatureLoss in the new notebook, but this is suboptimabl. Since the Loss is not needed for inference, I would like to be able to export my model without needing FeatureLoss to load it back.
Is it possible to do so in fastai, and if yes, how should I proceed?

4 Likes