Update code example in readme to use `device` variable instead of hard coded `cuda`

#4

The code example uses device for model creation but then uses a hard coded cuda for inputs. Updated to use device for inputs as well.

Xenova changed pull request status to merged

Sign up or log in to comment