device of noise models #634
Description
There are some inconsistencies when setting the device of the parameters of a noise model. The issue is that the device is decided by the random seed torch.Generator input of the noise model, instead of the device of the input parameter.
import deepinv as dinv
import torch
physics = dinv.physics.Denoising(noise_model=dinv.physics.GaussianNoise(sigma=0.01), device='cuda')
print(physics.noise_model.sigma.device)
# cpu while I would expect cuda
physics = dinv.physics.Denoising(noise_model=dinv.physics.GaussianNoise(sigma=torch.tensor(0.01, device='cuda')), device='cuda')
print(physics.noise_model.sigma.device)
# cpu while I would expect cuda
generator = torch.Generator(device='cuda')
physics = dinv.physics.Denoising(noise_model=dinv.physics.GaussianNoise(sigma=0.01, rng=generator), device='cuda')
print(physics.noise_model.sigma.device)
# cuda - OK
I believe that most users don't use the random seed, and thus can have unexpected device errors that are hard to spot. I propose that the device is decided by the input parameter instead of the generator.