code: def test_jit(self, device, dtype):
      inputs = torch.rand(2, 3, 30, 30, dtype=dtype, device=device)
  
      op = StraightThroughEstimator(torch.nn.MaxPool2d(3), grad_fn=None)
      op_script = torch.jit.script(op)
      actual = op_script(inputs)
      expected = op(inputs)
      assert_close(actual, expected)
  
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 4 for calling the TestSTE obj's underlying member method code with example usage: obj.test_jit(device, dtype) without return types