Error when includes any torchvision transforms in the pipeline | Cannot determine Numba type of <class 'ffcv.transforms.module.ModuleWrapper'>

[anasserhussien]

I have the following pipeline, the training loop runs properly without having torchvision.transforms.ColorJitter(.4,.4,.4) in the image pipeline, but once I include any transformation from torchvision.transforms it gives the following error.

CIFAR_MEAN = [125.307, 122.961, 113.8575]
CIFAR_STD = [51.5865, 50.847, 51.255]

BATCH_SIZE = 512
GPU_IDX = 0
device = torch.device(f"cuda:{GPU_IDX}") if torch.cuda.is_available() else torch.device("cpu")

loaders = {}
for name in ['train', 'test']:
    label_pipeline: List[Operation] = [IntDecoder(), ToTensor(), ToDevice(device), Squeeze()]
    image_pipeline: List[Operation] = [
        SimpleRGBImageDecoder(),
        RandomHorizontalFlip(),
        torchvision.transforms.ColorJitter(.4,.4,.4),
        RandomTranslate(padding=2),
        ToTensor(),
        ToDevice(device, non_blocking=True),
        ToTorchImage(),
        Convert(torch.float16),
        torchvision.transforms.Normalize(CIFAR_MEAN, CIFAR_STD), # Normalize using image statistics
    ]
    # Create loaders
    loaders[name] = Loader(f'/tmp/dataset_{name}.beton',
                            batch_size=BATCH_SIZE,
                            num_workers=8,
                            order=OrderOption.RANDOM,
                            drop_last=(name == 'train'),
                            pipelines={'image': image_pipeline,
                                       'label': label_pipeline})

Error

raceback (most recent call last):
  File "/home/ahmed/anaconda3/envs/ffcv/lib/python3.9/threading.py", line 980, in _bootstrap_inner
    self.run()
  File "/home/ahmed/anaconda3/envs/ffcv/lib/python3.9/site-packages/ffcv/loader/epoch_iterator.py", line 79, in run
    result = self.run_pipeline(b_ix, ixes, slot, events[slot])
  File "/home/ahmed/anaconda3/envs/ffcv/lib/python3.9/site-packages/ffcv/loader/epoch_iterator.py", line 133, in run_pipeline
    result = code(*args)
  File "/home/ahmed/anaconda3/envs/ffcv/lib/python3.9/site-packages/numba/core/dispatcher.py", line 468, in _compile_for_args
    error_rewrite(e, 'typing')
  File "/home/ahmed/anaconda3/envs/ffcv/lib/python3.9/site-packages/numba/core/dispatcher.py", line 409, in error_rewrite
    raise e.with_traceback(None)
numba.core.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Failed in nopython mode pipeline (step: nopython frontend)
Untyped global name 'self': Cannot determine Numba type of <class 'ffcv.transforms.module.ModuleWrapper'>

File "../anaconda3/envs/ffcv/lib/python3.9/site-packages/ffcv/transforms/module.py", line 25:
        def apply_module(inp, _):
            res = self.module(inp)
            ^

During: resolving callee type: type(CPUDispatcher(<function ModuleWrapper.generate_code.<locals>.apply_module at 0x7fef5a591310>))
During: typing of call at  (2)

During: resolving callee type: type(CPUDispatcher(<function ModuleWrapper.generate_code.<locals>.apply_module at 0x7fef5a591310>))
During: typing of call at  (2)

File "/home/ahmed/notebooks", line 2:
<source missing, REPL/exec in use?>

[pomonam]

Were you able to solve the issue? I am facing a same problem.

[anasserhussien]

Having ToTorchImage() in the pipeline before applying any trochvisoin transformation solved the issue!

[pomonam]

Great, that worked for me too. Thank you!