How to specify only a single non-zero-indexed GPU?

[NielsRogge]

Hi,

I'm on a shared device and would like to only leverage GPU with index = 1.

How to do this with Lightning?

I tried passing devices="1", devices=[1] but it's still leveraging the first 2 GPUs, as confirmed by this error:

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:1 and cuda:0! (when checking argument for argument index in method wrapper_CUDA__index_select)

My code is this notebook.