How to create a tensor that follows the model change device?
Post:
How to create a tensor in nn.Module, and his equipment changes with this model?
For example, I need a convolution kernel K, the parameters of which are given by me, and I don’t need to participate in training. But when I forward, I need him to perform several convolutions. In addition to this K, there are other convolutional layers.
If model().cuda(), this tensor is also on the GPU, if .cpu(), this tensor is also on the CPU.
Please use an online translation service before posting the question here, so that all users could help you.
I don’t think that’s entirely true, as my posted translation contains more details than the initial post, which might be helpful to answer the question.
@ChenzhouWeiYu to register parameters (such as conv weights) in the forward method of your model, you can either use self.register_parameter or assign the nn.Parameter to an internal attribute via self.param = nn.Parameter(...).