NameError: name 'xmp' is not defined

[pankaj-kvhld]

Working on Kaggle and trying to use TPU with fastai.

from fastai_xla_extensions.all import *

path = untar_data(URLs.MNIST_TINY)
datablock = DataBlock(
    blocks=(ImageBlock,CategoryBlock),
    get_items=get_image_files,
    splitter=GrandparentSplitter(),
    get_y=parent_label,
    item_tfms=Resize(28),
    batch_tfms=[Normalize.from_stats(*imagenet_stats)]
)
dls = datablock.dataloaders(path, bs=8) 
learner = cnn_learner(dls, resnet18, metrics=accuracy, concat_pool=False)
learner.xla_lr_find()

Throws the following error:

NameError                                 Traceback (most recent call last)
Cell In[7], line 1
----> 1 learner.xla_fit_one_cycle(5,lr_max=slice(1e-4,0.02))

File /usr/local/lib/python3.8/site-packages/fastai_xla_extensions/multi_core/learner.py:261, in xla_fit_one_cycle(self, n_epoch, num_cores, start_method, master_cbs, **kwargs)
    259 """call fit_one_cycle in a multicore tpu environment"""
    260 ctrl_args = self.pre_xla_fit()
--> 261 learner_args, add_args = self.pack_learner_args()
    263 fit_args = prep_fit_args(n_epoch, master_cbs, **kwargs)
    265 xmp.spawn(xla_run_method,
    266           args=(Learner.fit_one_cycle, learner_args, add_args, fit_args, ctrl_args),
    267           nprocs=num_cores,
    268           start_method=start_method)

File /usr/local/lib/python3.8/site-packages/fastai_xla_extensions/multi_core/learner.py:153, in pack_learner_args(self)
    151 "pack learner args into dict to pass to spawned process"
    152 learner_args = {**self.__stored_args__}
--> 153 learner_args['wrapped_model'] =  xmp.MpModelWrapper(self.model)
    154 learner_args['base_dls'] = self.dls
    155 # fetch only cbs not in defaults

NameError: name 'xmp' is not defined

I have no idea how to resolve this, any help will be greatly appreciated.
Thank you.