usage snippet

[yossilevii100]

Hi, thanks for the provided code. The usage code snippet yields errors if I ran it as is. The corrections are easy, you can add the following three lines.

from PIL import Image
import torch
device = "cuda" if torch.cuda.is_available() else "cpu"

Hope it helps and help others.