In some cases odeint works on GPU but fails on CPU

[talesa]

Hi!
I've came across an example that runs on the GPU but doesn't run on the CPU.
The cause seems to be x.norm() returning a nan when on CPU and a correct result on a GPU.
(https://github.com/rtqichen/torchdiffeq/blob/master/torchdiffeq/_impl/misc.py#L74)
This is happening when setting initial step at
https://github.com/rtqichen/torchdiffeq/blob/master/torchdiffeq/_impl/misc.py#L126

Don't have time to follow through on this at the moment, leaving it here for now.

[programmin1]

How did you set it to use GPU or CPU?

[mancunian1792]

I guess the adjoint method doesn't support GPU. odeint_adjoint is referred as odeint in the mnist example if you set the adjoint parameter to be true.