DataBlock resulting in RuntimeError: CUDA error: using a segmentation data

Part 1 (2020)
1

I am trying to build Datablock a kaggle dataset.
salt-identification

salt_block = DataBlock(blocks=(ImageBlock, MaskBlock),
                   get_items= images_list,
                   get_y = get_y_fn,
                   batch_tfms=[*aug_transforms(size=size), Normalize.from_stats(*imagenet_stats)])

Getting the following error:

RuntimeError: CUDA error: device-side assert triggered
CUDA kernel errors might be asynchronously reported at some other API call,so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.

It runs fine is i do not use the Normalise operation

salt_block = DataBlock(blocks=(ImageBlock, MaskBlock),
                   get_items= images_list,
                   get_y = get_y_fn,
                   batch_tfms=[*aug_transforms(size=size)])

Is there a way to get this working with Normalise operation.

Link to notebook: