Skip to content

Instantly share code, notes, and snippets.

@elumixor

elumixor/TRPO_flat_grad.py

Created May 24, 2020 23:31
Show Gist options
  • Save elumixor/78eabca2d952e3b3cef362b9ae2e2527 to your computer and use it in GitHub Desktop.
Save elumixor/78eabca2d952e3b3cef362b9ae2e2527 to your computer and use it in GitHub Desktop.
Download ZIP
Raw
TRPO_flat_grad.py
def flat_grad(y, x, retain_graph=False, create_graph=False):
if create_graph:
retain_graph = True
g = torch.autograd.grad(y, x, retain_graph=retain_graph, create_graph=create_graph)
g = torch.cat([t.view(-1) for t in g])
return g
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment