So it looks to me like in the Chainer code, the variables in k and g are treated as constant and are not backpropagated through; one way to do that in the PyTorch code would be to use param.grad.detach() rather than param.grad.clone().