Incorrect output device

Line 269 in 8e7931c

auto options = torch::TensorOptions().device(torch::kCUDA, 0);

Outputs are forced onto cuda:0 regardless of original device, breaking ddp. All the cuda functions seem to have this issue.

I've fixed it, will send a PR soon.

Thank you very much, you are completely right.
However, you changed other things in your PR: you added "nnMorpho." to the imports in several places. I imagine that you have import problems but this may be because you set "nnMorpho" as the project folder instead of "nnMorpho/nnMorpho". I usually put "nnMorpho/nnMorpho" as my Python project main folder because the "nnMorpho" folder contains other things for setting up the needed stuff for uploading to PyPi (like setup.py, README.md, etc...).
May you undo this adds to your PR such that I can only accept the ones about the CUDA device?

Installing with pip the package doesn't work at all without the import changes... (from pypi or from this repo)
Currently with those changes I'm able to use the library with pip install -e . in the repo and the same in CUDA

If there is some other fix coming I can rescind the import changes...

Ok. I merged your pull request.