Solved by ptrblck in post #2 
              
              
                You should register the model parameters as nn.Parameter. Otherwise the tensors won’t be properly registered. 
Alternatively, you could call register_parameter on the tensors. 
class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.l1 = nn.Parameter(torch.…