Expected cuda:0 but got CPU error

Did you move your Model as well to the device using Model.to(device)?