code: def the_test_with_annotations(self, relu):
      torch.manual_seed(100)
      model = TorchModel(relu)
      dummy_input = torch.rand(3, 1, 28, 28)
  
      config_list = [{"op_types": ["Conv2d"], "sparsity": 0.5}]
      pruner = L1NormPruner(model=model, config_list=config_list)
      _, masks = pruner.compress()
      pruner.unwrap_model()  # unwrap all modules to normal state
  
      masks["relu1"] = {
          "_input_input": torch.ones((8, 20, 24, 24)),
          "_output_0": torch.ones((8, 20, 24, 24)),
      }
      masks["conv1"]["_output_0"] = torch.ones((8, 20, 24, 24))
  
      traced_model = concrete_trace(
          model,
          {"x": dummy_input},
          leaf_module=(WithAnno1, WithAnno2, WithAnno3),
      )
      ModelSpeedup(traced_model, (dummy_input,), masks).speedup_model()
      traced_model(dummy_input)
  
      print("model before speedup", repr(model))
      # 125.49 M, 0.85M, 93.29, 1.1012
      flops, params, _ = count_flops_params(
          model, dummy_input, verbose=False
      )
      print(
          f"Pretrained model FLOPs {flops/1e6:.2f} M, #Params: {params/1e6:.2f}M"
      )
  
      print("model after speedup", repr(traced_model))
      flops, params, _ = count_flops_params(
          traced_model, dummy_input, verbose=False
      )
      print(
          f"Pruned model FLOPs {flops/1e6:.2f} M, #Params: {params/1e6:.2f}M"
      )
  
imports: import unittest
  import functools
  import torch
  import torch.nn.functional
  from nni.common.concrete_trace_utils import concrete_trace
  from nni.compression.pruning import L1NormPruner
  from nni.compression.speedup import ModelSpeedup
  from nni.compression.utils.counter import count_flops_params
label_desc: Use image node_id 1 for calling the InitMaskTestCase obj's underlying member method code with example usage: obj.the_test_with_annotations(relu) without return types