Hi @hungsing92 - unfortunately this is not possible at the moment, and it would be a big effort to make it so. All of the minimizer ops are wrapped with torch.no_grad (excluding only gradient computations). In addition to allowing grad, we'd also need to change all torch.autograd.grad calls to be passed create_graph=True, which adds a lot of overhead for the situations where this functionality is not needed.
Hi @hungsing92 - unfortunately this is not possible at the moment, and it would be a big effort to make it so. All of the minimizer ops are wrapped with
torch.no_grad(excluding only gradient computations). In addition to allowing grad, we'd also need to change alltorch.autograd.gradcalls to be passedcreate_graph=True, which adds a lot of overhead for the situations where this functionality is not needed.