chore(optim): wrap torch.autograd.grad() with torch.enable_grad() context

chore(optim): wrap torch.autograd.grad() with torch.enable_grad() context #1695