Cannot pickle a learner with FocalLossFlat loss function

Am getting the following error while trying to say model with FocalLossFlat() loss function. Same model is saved without problems if I use CrossEntropyFlat() instead.

learn = Learner(dls, model, loss_func=FocalLossFlat(gamma=gamma))
learn.fit_one_cycle(1, lr, wd=wd)
TypeError                                 Traceback (most recent call last)
Cell In[85], line 4
      1 # Save model
      2 # Export pickles the entire learner object. Including the dataloaders, loss function, optimizer, augmentations 
      3 # or transforms, and all callbacks. But will not point to any data
----> 4 learn.export(modelpath/"atm_recommender_export.pkl")

File ~/Workbench/fastai/fastai/fastai/, in export(self, fname, pickle_module, pickle_protocol)
    430 with warnings.catch_warnings():
    431     #To avoid the warning that come from PyTorch about model not being checked
    432     warnings.simplefilter("ignore")
--> 433, self.path/fname, pickle_module=pickle_module, pickle_protocol=pickle_protocol)
    434 self.create_opt()
    435 if state is not None: self.opt.load_state_dict(state)

File ~/.miniconda/envs/fastai/lib/python3.9/site-packages/torch/, in save(obj, f, pickle_module, pickle_protocol, _use_new_zipfile_serialization)
    421 if _use_new_zipfile_serialization:
    422     with _open_zipfile_writer(f) as opened_zipfile:
--> 423         _save(obj, opened_zipfile, pickle_module, pickle_protocol)
    424         return
    425 else:

File ~/.miniconda/envs/fastai/lib/python3.9/site-packages/torch/, in _save(obj, zip_file, pickle_module, pickle_protocol)
    633 pickler = pickle_module.Pickler(data_buf, protocol=pickle_protocol)
    634 pickler.persistent_id = persistent_id
--> 635 pickler.dump(obj)
    636 data_value = data_buf.getvalue()
    637 zip_file.write_record('data.pkl', data_value, len(data_value))

TypeError: cannot pickle 'code' object