model.cuda() won’t affect it, unless it has be reassigned. However, if you call model.cuda() and then forward a CUDA input, input.new will allocate a CUDA tensor, so the types will always match. I find that solution simpler and more robust than what you proposed. Doesn’t it work for you?