code: def test_one_hot_argmax_ste(self):
      for dtype in (torch.float, torch.double):
          X = torch.rand(5, 4, device=self.device, dtype=dtype)
          expected_discretized_X = one_hot(
              X.argmax(dim=-1), num_classes=X.shape[-1]
          ).to(X)
          discretized_X = OneHotArgmaxSTE.apply(X)
          # test forward
          self.assertTrue(
              torch.equal(expected_discretized_X, discretized_X)
          )
          # test backward
          X = X.requires_grad_(True)
          output = OneHotArgmaxSTE.apply(X)
          # sample some weights to checked that gradients are passed
          # as intended
          w = torch.rand_like(X)
          (w * output).sum().backward()
          self.assertTrue(torch.equal(w, X.grad))
  
file_path: botorch/test/utils/test_rounding.py
imports: import torch
  from botorch.utils.rounding import approximate_round, IdentitySTEFunction, OneHotArgmaxSTE, RoundSTE
  from botorch.utils.testing import BotorchTestCase
  from torch.nn.functional import one_hot
label_desc: Use image node_id 1 for calling the TestOneHotArgmaxSTE obj's underlying member method code with example usage: obj.test_one_hot_argmax_ste() without return types