Hi, I want to implement some network modules like Conv2D, which could be utilized in different models and different devices. Similarly to Conv2D, there are also some variables like weights/biases.
However, it seems that when I use torch.Tensor, it usually assumes or requires the device argument, which could not be decided. I have tried the .to function, but it doesn’t work for no reason (no warnings). The variables should be defined in init function instead of in the forward function for each forward pass I think. Therefore, the input.device can not be utilized.
Then, my question is: how to new/construct variables in the subclass of nn.Module that are agnostic to devices?
It could be simple but really confusing. Can anyone help me? Thank you!