Fix errors on CUDA device

[X3NNY]

Set parameters on the correct device.

[AlessandroFlati]

I made the exact same changes to make CUDA work. Nonetheless, it seems slightly less performance on "science" tasks, probably because of the time taken to port memory to cpu for some critical operations (the ones you found out needed a .to("cpu")).
@KindXiaoming would you mind explaining why is that necessary? I'm working hard on this repo believing in its potentiality, but I didn't have time to look at all the low level details.

[KindXiaoming]

Hi AlessandroFlati, do you mean why there are some operations are needed to perform on CPUs? In some cases: (1) I want to use sklearn (actually I never checked if sklearn can support GPU! so I just went with CPU). For example, sklearn linear regression will be called if you want to extract symbolic equations (which is needed if you want interpretability for science). (2) it seems to me the GPU version of torch.lstsq is more unstable than CPU version. Stable in the sense that sometimes GPU may return NaN and CPU doesn't. For example, torch.lstsq is used in curve2coef hence in initialize_grid_from_another_model (which is needed if you want to do grid extension to aim for accuracy for science).

[AlessandroFlati]

Then I think I can help you on that. I'll open a fork and try to prepare a PR demonstrating those 2 things can be performed in CUDA with dedicated examples. Maybe I should open another Issue for that, but NaN are also easily propagated if one sets the last activation to be exp for large numbers on dataset (~1e6 is enough) and I was wondering why. I'll keep an eye on that too.

[mw66]

@KindXiaoming

I think this fix deserve a new release (right now the version is v0.0.2).