Skip to content

Commit

Permalink
Merge remote-tracking branch 'origin/torch.Tensor.masked_select' into…
Browse files Browse the repository at this point in the history
… torch.Tensor.masked_select
  • Loading branch information
Jin Wang committed Jun 30, 2024
2 parents c27fda5 + 4eb1d2e commit 10f4560
Showing 1 changed file with 1 addition and 3 deletions.
4 changes: 1 addition & 3 deletions ivy/functional/frontends/torch/tensor.py
Original file line number Diff line number Diff line change
Expand Up @@ -1072,9 +1072,7 @@ def masked_fill_(self, mask, value):
return self

def masked_select(self, mask):
x, mask = torch_frontend.broadcast_tensors(self, mask)
out = x.reshape(-1).index_select(0, mask.reshape(-1).nonzero().flatten())
return torch_frontend.tensor(out)
return torch_frontend.masked_select(self, mask)

@with_unsupported_dtypes({"2.2 and below": ("float16", "bfloat16")}, "torch")
def index_add_(self, dim, index, source, *, alpha=1):
Expand Down

0 comments on commit 10f4560

Please sign in to comment.