nnabla
nnabla copied to clipboard
Add to_grad option to F.assign
Hi, @TE-AkioHayakawa san , @TE-TakuyaNarihira san.
I've added to_grad option to F.assign function. This will be useful to build original your own gradient manipulation with use of nn.grad.
x = nn.Variable((2, 3), need_grad=True)
y = nn.Variable((2, 3))
l = F.mean(F.squared_error(x, y))
# backward
grad = nn.grad([l], [x], bind_grad_output=True)[0]
# clip gradients by norm
clipped_grad = F.clip_by_norm(grad, 0.5)
# apply clipped gradients to x
assign = F.assign(x, clipped_grad, to_grad=True)
assign.forward(clear_buffer=True)