usage snippet

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.
