Can't be in JIT mode and on the GPU

[samuelstevens]

I would like to do as much work on the CPU to avoid taking up GPU memory, but I get

AssertionError: Can't be in JIT mode and on the GPU

How can I indicate to the loader that some operations need JIT and some are not?

Dataloader code (I don't think the writer file is necessary).

return ffcv.loader.Loader(
    ffcv_path,
    batch_size=batch_size,
    num_workers=16,
    order=ffcv.loader.OrderOption.RANDOM,
    os_cache=True,
    distributed=True,
    drop_last=False,
    pipelines={
        "image": [
            ffcv.fields.decoders.SimpleRGBImageDecoder(),
            ffcv.transforms.NormalizeImage(
                self.mean.numpy(), self.std.numpy(), np.float32
            ),
            ffcv.transforms.ToTensor(),
            ffcv.transforms.ToTorchImage(),
            ffcv.transforms.ToDevice(accelerator.device),
        ],
        "label": [
            ffcv.fields.decoders.NDArrayDecoder(),
            ffcv.transforms.ToTensor(),
            ffcv.transforms.Convert(torch.int64),
            ffcv.transforms.ToDevice(accelerator.device),
        ],
    },
)

[samuelstevens]

I tried fixing it by switching the transforms in "label", specifically, to:

"label": [
    ffcv.fields.decoders.NDArrayDecoder(),
    ffcv.transforms.Convert(np.int64),
    ffcv.transforms.ToTensor(),
    ffcv.transforms.ToDevice(accelerator.device),
]

But now I get: Untyped global name 'self': Cannot determine Numba type of <class 'ffcv.transforms.ops.Convert'>

[samuelstevens]

Fixed with:

"label": [
    ffcv.fields.decoders.NDArrayDecoder(),    
    ffcv.transforms.ToTensor(),
    ffcv.transforms.Convert(torch.int64),
    ffcv.transforms.ToDevice(accelerator.device),
]

Unclear to me why this makes a difference but there are no more errors!

[niniack]

Hi @samuelstevens , I'm running into the same initial error. Just wanted to ask you, since it is not obvious to me, what is the difference between the first and the last label pipeline? It seems you are using the same decoders and transforms in the same order.

[samuelstevens]

I switched the order of ToTensor and Convert(torch.int64)