How to get TensorMask to work?

I am trying to follow Vision core notebook ( and am running into a problem. I get the error that TensorMask used in checking the type is not defined. I searched around the Fastai github and found the TensorMask class defined in torch_core. I imported it from there but it is still giving me the same error. What I might be doing wrong?

from fastai.torch_core import TensorMask

class BypassNewMeta(type):
    "Metaclass: casts `x` to this class, initializing with `_new_meta` if available"
    def __call__(cls, x, *args, **kwargs):
        if hasattr(cls, '_new_meta'): x = cls._new_meta(x, *args, **kwargs)
        if cls!=x.__class__: x.__class__ = cls
        return x
class PILBase(Image, metaclass=BypassNewMeta):
    _show_args = {'cmap':'viridis'}
    _open_args = {'mode': 'RGB'}
    def create(cls, fn, **kwargs)->None:
        "Open an `Image` from path `fn`"
        if isinstance(fn,TensorImage): fn = fn.permute(1,2,0).type(torch.uint8)
        if isinstance(fn, TensorMask): fn = fn.type(torch.uint8)
        if isinstance(fn,Tensor): fn = fn.numpy()
        if isinstance(fn,ndarray): return cls(Image.fromarray(fn))
        if isinstance(fn,bytes): fn = io.BytesIO(fn)
        return cls(load_image(fn, **merge(cls._open_args, kwargs)))

    def show(self, ctx=None, **kwargs):
        "Show image using `merge(self._show_args, kwargs)`"
        return show_image(self, ctx=ctx, **merge(self._show_args, kwargs))

    def __repr__(self): return f'{self.__class__.__name__} mode={self.mode} size={"x".join([str(d) for d in self.size])}'

class PILImage(PILBase): pass

ex_image = PILImage.create(ex_image) #the error appears here

I hope that code will be formatted well