Solved by mulholio in post #2 
              
              
                Found the fix. The tensor returned by init_params needs to be wrapped in .requires_grad_(), not the individual tensors within it. 
def init_params(size, variance=1.0):
    return (torch.randn(size, dtype=torch.float)*variance).requires_grad_()