How to run MONAI on GPU?

[ghost]

Hello,
I've been trying to vizualise data with plt, but my code crashes because theres not enough RAM.
The line that makes the code crash is check_data = first(check_loader)
I'm coding on Google Colab because I don't have a nvidia GPU, and I've already changes the execution type to GPU, but I can't get it to work...
Does anyone know how to fix this ?

Here is the part of the code that I struggle to run:
data_dir = "/content/drive/MyDrive/EPFL/Projet Bachelor/Zebrafish_project/Datasets/"
split_json = "dataset.json"

datasets = data_dir + split_json
train_datalist = load_decathlon_datalist(datasets, True, "training")
#val_datalist = load_decathlon_datalist(datasets, True, "validation")
print(train_datalist)

check_ds = Dataset(data = train_datalist, transform=train_transforms)
check_loader = DataLoader(check_ds, batch_size=1)
check_data = first(check_loader)
check_data = check_data.to(device)
image, label = (check_data["image"][0][0], check_data["label"][0][0])
print(f"image shape: {image.shape}, label shape: {label.shape}")
plt.figure("check", (12, 6))
plt.subplot(1, 2, 1)
plt.title("image")
plt.imshow(image[:, :, 80], cmap="gray")
plt.subplot(1, 2, 2)
plt.title("label")
plt.imshow(label[:, :, 80])
plt.show()

Thank you.

[Lucas-rbnt]

There's no reason why you can't load an image on the CPU, as we generally keep the GPU for training models. Besides, it seems to me that it's impossible to display images from the gpu (at least straightforward plt.imshow and plt.matshow requires input to be passed to cpu first).

Can you provide more information? What error did you find? Can you copy/paste what is being returned? What dataset.json file are you using, etc.?