If you want it to run on the GPU, you can just do
inputs.to('cuda:0') model.to('cuda:0')
· Sign up or log in to comment