code: def test_function(self, device, dtype):
      inpt = torch.randn(
          4, requires_grad=True, device=device, dtype=dtype
      )
      output = torch.sign(inpt)
      loss = output.mean()
      loss.backward()
      assert_close(
          inpt.grad,
          torch.tensor(
              [0.0, 0.0, 0.0, 0.0], device=device, dtype=dtype
          ),
      )
  
      out_est = STEFunction.apply(inpt, output, F.hardtanh)
      loss = out_est.mean()
      loss.backward()
      assert_close(
          inpt.grad,
          torch.tensor(
              [0.2500, 0.2500, 0.2500, 0.2500],
              device=device,
              dtype=dtype,
          ),
      )
  
      out_est = STEFunction.apply(inpt, output, None)
      loss = out_est.mean()
      loss.backward()
      assert_close(
          inpt.grad,
          torch.tensor(
              [0.5000, 0.5000, 0.5000, 0.5000],
              device=device,
              dtype=dtype,
          ),
      )
  
file_path: kornia/test/grad_estimator/test_ste.py
imports: import pytest
  import torch
  import torch.nn.functional
  from torch.testing import assert_close
  import kornia.augmentation
  from kornia.grad_estimator import STEFunction, StraightThroughEstimator
label_desc: Use image node_id 2 for calling the TestSTE obj's underlying member method code with example usage: obj.test_function(device, dtype) without return types