The error is raised if the input and parameters are on different devices, which could happen if you forget to push either one of them to the GPU.
Make sure to use model.to('cuda') as well as input = input.to('cuda') before executing the code.
PS: your current code is a bit hard to read so please format it by wrapping it into three backticks ```