default num_classes value for one_hot (#6182)

* num_classes=-1

If num_classes set to -1, the number of classes will be inferred as one greater than the largest class value in the input tensor.

* num_classes desc

comment to explain num_classes default and what that means.

* replacing ' with `
This commit is contained in:
Gabe Caldwell
2024-08-19 15:07:14 -04:00
committed by GitHub
parent 9328248610
commit bdd6325f31
2 changed files with 8 additions and 1 deletions

View File

@@ -2033,9 +2033,13 @@ class TestOps(unittest.TestCase):
data = [1, 2, 4]
helper_test_op([], lambda: torch.nn.functional.one_hot(torch.tensor(data), 6).type(torch.int32),
lambda: Tensor(data).one_hot(6), forward_only=True)
helper_test_op([], lambda: torch.nn.functional.one_hot(torch.tensor(data)).type(torch.int32),
lambda: Tensor(data).one_hot(), forward_only=True)
data = [[[1, 2, 3], [0, 3, 5]], [[1, 2, 3], [0, 3, 5]]]
helper_test_op([], lambda: torch.nn.functional.one_hot(torch.tensor(data), 8).type(torch.int32),
lambda: Tensor(data).one_hot(8), forward_only=True)
helper_test_op([], lambda: torch.nn.functional.one_hot(torch.tensor(data)).type(torch.int32),
lambda: Tensor(data).one_hot(), forward_only=True)
def test_masked_fill(self):
helper_test_op([(32,10)], lambda x: x.masked_fill((x>0.1).detach(), -math.inf))