Yes, well in my case I don’t really need to save this parameter into my state_dict. I just need that this parameter is moved between devices when call model.to(device) and that is not passed to the optimizer when calling model.parameters().
I could workaround this issue if I run in the forward of my nn.Module something like:
W = self.W.to(self.device)
Y = X@W # here X can be in GPU or CPU depending on the iteration 

But I would prefer not calling the method to.(self.device) on each forward and just handle the device of each of the parameters through a call to my_model.to()