Since you are initializing the tensors with requires_grad=True
, you should wrap them into nn.Parameter
, so that they will be properly registered in the state_dict
and will be automatically pushed to the device, if you call model.to(device)
.
Also, since you are storing these parameters in a list, use ParameterList
, as a plain Python list won’t register the parameters properly.
This should work:
class Model(nn.Module):
def __init__(self):
super(Model, self).__init__()
self.w_x = nn.ParameterList([nn.Parameter(torch.randn(NUM_INPUT)) for _ in range(NUM_HIDDEN)])
self.w_h = nn.ParameterList([nn.Parameter(torch.randn(NUM_INPUT)) for _ in range(NUM_HIDDEN)])
self.sel = nn.ParameterList([nn.Parameter(torch.zeros(NUM_INPUT)) for _ in range(NUM_HIDDEN)])
self.x = nn.ParameterList([nn.Parameter(torch.zeros(NUM_INPUT)) for _ in range(NUM_HIDDEN)])
self.h = nn.ParameterList([nn.Parameter(torch.zeros(NUM_INPUT)) for _ in range(NUM_HIDDEN)])
model = Model()
model.cuda()
print(model.parameters())