code: def test_custom_sparse_mm_op(self, input_shape=(4, 4)):
      class TestLayer(nn.Module):
          def __init__(self):
              super(TestLayer, self).__init__()
  
          def forward(self, x, y):
              x = torch.sparse.mm(x, y)
              return x
  
      model = TestLayer()
      input_data_x = torch.ones(input_shape)
      input_data_y = torch.ones(input_shape)
      input_data = [input_data_x, input_data_y]
      model.eval()
      torch_model = torch.jit.trace(model, (input_data_x, input_data_y))
      mlmodel = convert_to_mlmodel(torch_model, input_data)
  
      layers = mlmodel.get_spec().neuralNetwork.layers
      assert layers[-1].custom is not None, "Expecting a custom layer"
      assert (
          "SparseMatMul" == layers[-1].custom.className
      ), "Custom Layer class name mismatch"
      assert (
          not layers[-1].custom.parameters["transpose_x"].boolValue
      ), "Incorrect parameter value k"
      assert (
          not layers[-1].custom.parameters["transpose_y"].boolValue
      ), "Incorrect parameter value k"
      assert (
          layers[-1].custom.parameters["x_is_sparse"].boolValue
      ), "Incorrect parameter value k"
      assert (
          layers[-1].custom.parameters["y_is_sparse"].boolValue
      ), "Incorrect parameter value k"
  
file_path: coremltools/coremltools/converters/mil/frontend/torch/test/test_custom_ops.py
imports: import pytest
  import torch
  import torch.nn
  from coremltools.converters.mil.frontend.torch.ops import _get_inputs
  from coremltools.converters.mil.frontend.torch.ops import cosine_similarity
  from coremltools.converters.mil.frontend.torch.torch_op_registry import _TORCH_OPS_REGISTRY
  from coremltools.converters.mil.frontend.torch.torch_op_registry import register_torch_op
  from coremltools.converters.mil.mil import Builder
  from coremltools.converters.mil.mil import Operation, types
  from coremltools.converters.mil.mil.input_type import DefaultInputs, InputSpec, TensorInputType
  from coremltools.converters.mil.mil.ops.defs._op_reqs import register_op
  from .testing_utils import TorchBaseTest, convert_to_mlmodel
label_desc: Use image node_id 2 for calling the TestCustomOp obj's underlying member method code with example usage: obj.test_custom_sparse_mm_op(input_shape) and returns: x