Fix inconsistent devices for tensors

[hangtingchen]

Fix inconsistent devices for tensors

[jonashaag]

Looks good, can you add a test?

[mpariente]

IMO, we can merge without a test.

[jonashaag]

Ah, since we only have a single PyTorch device in CI, it's kind of pointless anyways. I was wondering why this isn't tested already, but we always only ever test with the default device in the first place. At some point it would be really great to add a second the device to CI and also run all of the tests on CPU and CUDA.

[hangtingchen]

The bug was found when I train the network on GPU, but cnt was placed on CPU in default. It is a good idea to add a second device to CI. Maybe this can be done in another PR.

[jonashaag]

It is a good idea to add a second device to CI. Maybe this can be done in another PR.

Unfortunately we don't have GPU support in CI :(

Will merge this fix, thanks for the patch.

[mpariente]

Yes, no GPU on CI. Maybe we can emulate another device? It might be too much work.
I don't know whether CircleCI or drone provide free GPU CI, we might want to have a look.