Skip to content

Commit

Permalink
Fix incorrect trace of post-accumulate grad hook on tensor with zero …
Browse files Browse the repository at this point in the history
…dimensions

pytorch#135226

Change-Id: Ic23442910ebe217a14a1762a78d0195b13d80111
  • Loading branch information
internal developer authored and aostrowski-hbn committed Oct 16, 2024
1 parent 78a7e04 commit b304b80
Showing 1 changed file with 2 additions and 2 deletions.
4 changes: 2 additions & 2 deletions torch/_dynamo/compiled_autograd.py
Original file line number Diff line number Diff line change
Expand Up @@ -199,13 +199,13 @@ def post_acc_grad_hook(self, input, hook_id):
assert isinstance(input, torch.Tensor)
assert self.hooks_proxy is not None
hook = self.hooks_proxy[hook_id] # type: ignore[index]
proxies = self.proxy_call_hook(
proxy = self.proxy_call_hook(
hook,
input,
)
with disable_proxy_modes_tracing():
input = [maybe_clone(input)]
self.bind_tensors_to_proxies(input, proxies)
self.bind_tensors_to_proxies(input, [proxy])
return input

# Note: [Compiled autograd and cudagraphs]
Expand Down

0 comments on commit b304b80

Please sign in to comment.