Skip to content

Commit

Permalink
[pre-commit.ci] auto fixes from pre-commit.com hooks
Browse files Browse the repository at this point in the history
for more information, see https://pre-commit.ci
  • Loading branch information
pre-commit-ci[bot] committed Apr 2, 2024
1 parent f31bf79 commit c7eef87
Show file tree
Hide file tree
Showing 2 changed files with 5 additions and 1 deletion.
4 changes: 3 additions & 1 deletion thunder/core/transforms.py
Expand Up @@ -1109,7 +1109,9 @@ def _sum_prim_grad(a: TensorProxy, /, dims: Sequence[int]) -> TensorProxy:


@torchctx
def _topk_prim_grad(a: TensorProxy, /, k: int, dim: None | int = None, largest: bool = True, sorted: bool = True, *, out=None):
def _topk_prim_grad(
a: TensorProxy, /, k: int, dim: None | int = None, largest: bool = True, sorted: bool = True, *, out=None
):
fwd = prims.topk(a, k, dim, largest, sorted, out=out)
val, idx = fwd

Expand Down
2 changes: 2 additions & 0 deletions thunder/tests/opinfos.py
Expand Up @@ -4771,6 +4771,8 @@ def topk_thunder_ref(*args, **kwargs):

def topk_torch_ref(*args, **kwargs):
return torch.topk(*args, **kwargs)[0]


# }


Expand Down

0 comments on commit c7eef87

Please sign in to comment.