Align MLPerf Tiny models with reference inference graphs (pytorch#22707)

### Summary
Align the MLPerf Tiny inference fixtures with the MLCommons reference
graphs. Correct SAME padding, enable convolution biases, and pool over
all 25x5 KWS features. Add optional classifier softmax and enable it in
the inference wrappers and Cortex-M model tests. Keep the fixed 16/32/64
ResNet stage widths and existing BatchNorm settings.

### Test plan
All eight Cortex-M model tests pass, including four Corstone-300 FVP
runs. All twelve TOSA model tests pass. Eight XNNPACK export checks and
sixteen Ethos-U55/U85 compilation checks pass; these do not include
target runtime execution. Formatting and Flake8 pass. Validation uses
the repository's pinned PyTorch 2.13 and torchao versions.

Authored with OpenAI Codex.

cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils
@Sebastian-Larsson @robell @robert-kalmar @JakeStevens

Author: rascani

backends/arm/test/models/test_ds_cnn.py

@@ -22,7 +22,7 @@
 _wrapper = DSCNNKWSModel()
 model = _wrapper.get_eager_model()
 # TODO: Remove once a pass decomposes large-stride AvgPool2d.
-# Replace AvgPool2d(24,5) with AdaptiveAvgPool2d(1) so the
+# Replace AvgPool2d(25,5) with AdaptiveAvgPool2d(1) so the
 # DecomposeAdaptiveAvgPool2dPass can break it into stride-1
 # pools that satisfy the Ethos-U55 stride <= 3 constraint.
 model.pool = nn.AdaptiveAvgPool2d(output_size=1)  # type: ignore[assignment]

backends/cortex_m/test/models/test_ds_cnn.py

@@ -10,14 +10,16 @@
 from executorch.examples.models.mlperf_tiny.ds_cnn import DSCNNKWS
 
 ops_before_transforms: dict[str, int] = {
+    "executorch_exir_dialects_edge__ops_aten__softmax_default": 1,
     "executorch_exir_dialects_edge__ops_aten_avg_pool2d_default": 1,
+    "executorch_exir_dialects_edge__ops_aten_constant_pad_nd_default": 1,
     "executorch_exir_dialects_edge__ops_aten_convolution_default": 9,
     "executorch_exir_dialects_edge__ops_aten_linear_default": 1,
     "executorch_exir_dialects_edge__ops_aten_relu_default": 9,
     "executorch_exir_dialects_edge__ops_aten_view_copy_default": 1,
     "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_channel_default": 18,
-    "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_tensor_default": 15,
-    "executorch_exir_dialects_edge__ops_quantized_decomposed_quantize_per_tensor_default": 13,
+    "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_tensor_default": 17,
+    "executorch_exir_dialects_edge__ops_quantized_decomposed_quantize_per_tensor_default": 15,
 }
 
 ops_after_transforms: dict[str, int] = {
@@ -29,11 +31,12 @@
     "executorch_exir_dialects_edge__ops_cortex_m_quantized_conv2d_default": 5,
     "executorch_exir_dialects_edge__ops_cortex_m_quantized_depthwise_conv2d_default": 4,
     "executorch_exir_dialects_edge__ops_cortex_m_quantized_linear_default": 1,
+    "executorch_exir_dialects_edge__ops_cortex_m_softmax_default": 1,
 }
 
 test_cases = {
     "ds_cnn": McuTestCase(
-        model=DSCNNKWS().eval(),
+        model=DSCNNKWS(apply_softmax=True).eval(),
         example_inputs=lambda: (
             (torch.rand(1, 1, 49, 10) * 2 - 1).to(memory_format=torch.channels_last),
         ),

backends/cortex_m/test/models/test_mobilenet_v1_025.py

@@ -10,14 +10,16 @@
 from executorch.examples.models.mlperf_tiny.mobilenet_v1_025 import MobileNetV1025
 
 ops_before_transforms: dict[str, int] = {
+    "executorch_exir_dialects_edge__ops_aten__softmax_default": 1,
     "executorch_exir_dialects_edge__ops_aten_avg_pool2d_default": 1,
+    "executorch_exir_dialects_edge__ops_aten_constant_pad_nd_default": 5,
     "executorch_exir_dialects_edge__ops_aten_convolution_default": 27,
     "executorch_exir_dialects_edge__ops_aten_linear_default": 1,
     "executorch_exir_dialects_edge__ops_aten_relu_default": 27,
     "executorch_exir_dialects_edge__ops_aten_view_copy_default": 1,
     "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_channel_default": 54,
-    "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_tensor_default": 33,
-    "executorch_exir_dialects_edge__ops_quantized_decomposed_quantize_per_tensor_default": 31,
+    "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_tensor_default": 39,
+    "executorch_exir_dialects_edge__ops_quantized_decomposed_quantize_per_tensor_default": 37,
 }
 
 ops_after_transforms: dict[str, int] = {
@@ -29,11 +31,12 @@
     "executorch_exir_dialects_edge__ops_cortex_m_quantized_conv2d_default": 14,
     "executorch_exir_dialects_edge__ops_cortex_m_quantized_depthwise_conv2d_default": 13,
     "executorch_exir_dialects_edge__ops_cortex_m_quantized_linear_default": 1,
+    "executorch_exir_dialects_edge__ops_cortex_m_softmax_default": 1,
 }
 
 test_cases = {
     "mobilenet_v1_025": McuTestCase(
-        model=MobileNetV1025().eval(),
+        model=MobileNetV1025(apply_softmax=True).eval(),
         example_inputs=lambda: (
             (torch.rand(1, 3, 96, 96) * 2 - 1).to(memory_format=torch.channels_last),
         ),

backends/cortex_m/test/models/test_resnet8.py

@@ -10,15 +10,17 @@
 from executorch.examples.models.mlperf_tiny.resnet8 import ResNet8
 
 ops_before_transforms: dict[str, int] = {
+    "executorch_exir_dialects_edge__ops_aten__softmax_default": 1,
     "executorch_exir_dialects_edge__ops_aten_add_Tensor": 3,
     "executorch_exir_dialects_edge__ops_aten_avg_pool2d_default": 1,
+    "executorch_exir_dialects_edge__ops_aten_constant_pad_nd_default": 2,
     "executorch_exir_dialects_edge__ops_aten_convolution_default": 9,
     "executorch_exir_dialects_edge__ops_aten_linear_default": 1,
     "executorch_exir_dialects_edge__ops_aten_relu_default": 7,
     "executorch_exir_dialects_edge__ops_aten_view_copy_default": 1,
-    "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_channel_default": 16,
-    "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_tensor_default": 21,
-    "executorch_exir_dialects_edge__ops_quantized_decomposed_quantize_per_tensor_default": 16,
+    "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_channel_default": 18,
+    "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_tensor_default": 24,
+    "executorch_exir_dialects_edge__ops_quantized_decomposed_quantize_per_tensor_default": 19,
 }
 
 ops_after_transforms: dict[str, int] = {
@@ -30,11 +32,12 @@
     "executorch_exir_dialects_edge__ops_cortex_m_quantized_avg_pool2d_default": 1,
     "executorch_exir_dialects_edge__ops_cortex_m_quantized_conv2d_default": 9,
     "executorch_exir_dialects_edge__ops_cortex_m_quantized_linear_default": 1,
+    "executorch_exir_dialects_edge__ops_cortex_m_softmax_default": 1,
 }
 
 test_cases = {
     "resnet8": McuTestCase(
-        model=ResNet8().eval(),
+        model=ResNet8(apply_softmax=True).eval(),
         example_inputs=lambda: (
             (torch.rand(1, 3, 32, 32) * 2 - 1).to(memory_format=torch.channels_last),
         ),

backends/nxp/tests/generic_tests/test_profiling.py

@@ -397,23 +397,23 @@ def test__resnet8(self, caplog, request):
         )
         neutron_map = extract_map_from_logs(caplog)
         assert neutron_map == {
-            0: (14, 15),  # Conv2DStandardV2 (Pad + Conv)
-            1: (17, 18),  # Conv2DStandardV1 (Pad + Conv)
-            2: (20,),  # Conv2DStandardV1
-            3: (21,),  # Add
-            4: (22,),  # GlobalBiasScale (Relu)
-            5: (28,),  # Conv2DStandardV1
-            6: (24, 25),  # Conv2DStandardV1 (Pad + Conv)
-            7: (27,),  # Conv2DStandardV1
-            8: (29,),  # Add
-            9: (30,),  # GlobalBiasScale (Relu)
-            10: (36,),  # Conv2DStandardV1
-            11: (32, 33),  # Conv2DStandardV1 (Pad + Conv)
-            12: (35,),  # Conv2DStandardV1
-            13: (37,),  # Add
-            14: (38,),  # GlobalBiasScale (Relu)
-            15: (39, 40),  # GlobalAvgPool (Mean + Reshape)
-            16: (41,),  # FullyConnected
+            0: (22, 23),  # Conv2DStandardV2 (Pad + Conv)
+            1: (24, 25),  # Conv2DStandardV1 (Pad + Conv)
+            2: (26,),  # Conv2DStandardV1
+            3: (27,),  # Add
+            4: (28,),  # GlobalBiasScale (Relu)
+            5: (33,),  # Conv2DStandardV1
+            6: (29, 30, 31),  # Conv2DStandardV1 (Pad + Conv)
+            7: (32,),  # Conv2DStandardV1
+            8: (34,),  # Add
+            9: (35,),  # GlobalBiasScale (Relu)
+            10: (40,),  # Conv2DStandardV1
+            11: (36, 37, 38),  # Conv2DStandardV1 (Pad + Conv)
+            12: (39,),  # Conv2DStandardV1
+            13: (41,),  # Add
+            14: (42,),  # GlobalBiasScale (Relu)
+            15: (43, 44),  # GlobalAvgPool (Mean + Reshape)
+            16: (45,),  # FullyConnected
             17: (),  # Neutron Dump
         }
 
@@ -434,34 +434,34 @@ def test__resnet8_batch_2(self, caplog, request):
         )
         neutron_map = extract_map_from_logs(caplog)
         assert neutron_map == {
-            0: (14, 15),  # Conv2DStandardV2 (1/2 batch item)
-            1: (14, 15),  # Conv2DStandardV2 (2/2 batch item)
-            2: (17, 18),  # Conv2DStandardV1 (1/2 batch item)
-            3: (17, 18),  # Conv2DStandardV1 (2/2 batch item)
-            4: (20,),  # Conv2DStandardV1 (1/2 batch item)
-            5: (20,),  # Conv2DStandardV1 (2/2 batch item)
-            6: (21,),  # Add
-            7: (22,),  # GlobalBiasScale (Relu)
-            8: (28,),  # Conv2DStandardV1 (1/2 batch item)
-            9: (28,),  # Conv2DStandardV1 (2/2 batch item)
-            10: (24, 25),  # Conv2DStandardV1 (1/2 batch item)
-            11: (24, 25),  # Conv2DStandardV1 (2/2 batch item)
-            12: (27,),  # Conv2DStandardV1 (1/2 batch item)
-            13: (27,),  # Conv2DStandardV1 (2/2 batch item)
-            14: (29,),  # Add
-            15: (30,),  # GlobalBiasScale (Relu)
-            16: (36,),  # Conv2DStandardV1 (1/2 batch item)
-            17: (36,),  # Conv2DStandardV1 (2/2 batch item)
-            18: (32, 33),  # Conv2DStandardV1 (1/2 batch item)
-            19: (32, 33),  # Conv2DStandardV1 (2/2 batch item)
-            20: (35,),  # Conv2DStandardV1 (1/2 batch item)
-            21: (35,),  # Conv2DStandardV1 (2/2 batch item)
-            22: (37,),  # Add
-            23: (38,),  # GlobalBiasScale (Relu)
-            24: (39,),  # GlobalAvgPool (1/2 batch item)
-            25: (39,),  # GlobalAvgPool (2/2 batch item)
-            26: (41,),  # Conv2DPointwise (FullyConnected)
-            27: (41,),  # Slice (FullyConnected)
+            0: (22, 23),  # Conv2DStandardV2 (1/2 batch item)
+            1: (22, 23),  # Conv2DStandardV2 (2/2 batch item)
+            2: (24, 25),  # Conv2DStandardV1 (1/2 batch item)
+            3: (24, 25),  # Conv2DStandardV1 (2/2 batch item)
+            4: (26,),  # Conv2DStandardV1 (1/2 batch item)
+            5: (26,),  # Conv2DStandardV1 (2/2 batch item)
+            6: (27,),  # Add
+            7: (28,),  # GlobalBiasScale (Relu)
+            8: (33,),  # Conv2DStandardV1 (1/2 batch item)
+            9: (33,),  # Conv2DStandardV1 (2/2 batch item)
+            10: (29, 30, 31),  # Conv2DStandardV1 (1/2 batch item)
+            11: (29, 30, 31),  # Conv2DStandardV1 (2/2 batch item)
+            12: (32,),  # Conv2DStandardV1 (1/2 batch item)
+            13: (32,),  # Conv2DStandardV1 (2/2 batch item)
+            14: (34,),  # Add
+            15: (35,),  # GlobalBiasScale (Relu)
+            16: (40,),  # Conv2DStandardV1 (1/2 batch item)
+            17: (40,),  # Conv2DStandardV1 (2/2 batch item)
+            18: (36, 37, 38),  # Conv2DStandardV1 (1/2 batch item)
+            19: (36, 37, 38),  # Conv2DStandardV1 (2/2 batch item)
+            20: (39,),  # Conv2DStandardV1 (1/2 batch item)
+            21: (39,),  # Conv2DStandardV1 (2/2 batch item)
+            22: (41,),  # Add
+            23: (42,),  # GlobalBiasScale (Relu)
+            24: (43,),  # GlobalAvgPool (1/2 batch item)
+            25: (43,),  # GlobalAvgPool (2/2 batch item)
+            26: (45,),  # Conv2DPointwise (FullyConnected)
+            27: (45,),  # Slice (FullyConnected)
             28: (),  # Neutron Dump
         }
         inspector_check(get_test_name(request))
@@ -483,18 +483,18 @@ def test__ds_cnn(self, caplog, request):
         )
         neutron_map = extract_map_from_logs(caplog)
         assert neutron_map == {
-            0: (14, 15),  # Pad (Conv + Relu)
-            1: (14, 15),  # Conv2DStandardV2 (Conv + Relu)
-            2: (18, 19),  # Conv2DDepthwiseV1 (DepthwiseConv + Relu)
-            3: (21, 22),  # Conv2DPointwise (Conv + Relu)
-            4: (24, 25),  # Conv2DDepthwiseV1 (DepthwiseConv + Relu)
-            5: (27, 28),  # Conv2DDepthwiseV1 (Conv + Relu)
-            6: (30, 31),  # Conv2DDepthwiseV1 (DepthwiseConv + Relu)
-            7: (33, 34),  # Conv2DPointwise (Conv + Relu)
-            8: (36, 37),  # Conv2DDepthwiseV1 (DepthwiseConv + Relu)
-            9: (39, 40),  # Conv2DPointwise  (Conv + Relu)
-            10: (42,),  # Conv2DDepthwiseDense (AvgPool)
-            11: (44,),  # FullyConnected
+            0: (22, 23, 24),  # Pad (Pad + Conv + Relu)
+            1: (22, 23, 24),  # Conv2DStandardV2 (Pad + Conv + Relu)
+            2: (26, 27),  # Conv2DDepthwiseV1 (DepthwiseConv + Relu)
+            3: (28, 29),  # Conv2DPointwise (Conv + Relu)
+            4: (30, 31),  # Conv2DDepthwiseV1 (DepthwiseConv + Relu)
+            5: (32, 33),  # Conv2DDepthwiseV1 (Conv + Relu)
+            6: (34, 35),  # Conv2DDepthwiseV1 (DepthwiseConv + Relu)
+            7: (36, 37),  # Conv2DPointwise (Conv + Relu)
+            8: (38, 39),  # Conv2DDepthwiseV1 (DepthwiseConv + Relu)
+            9: (40, 41),  # Conv2DPointwise  (Conv + Relu)
+            10: (43, 44),  # Conv2DDepthwiseDense (AvgPool + Reshape)
+            11: (45,),  # FullyConnected
             12: (),  # Neutron Dump
         }
 
@@ -515,35 +515,35 @@ def test__mobilenet_v1_025(self, caplog, request):
         )
         neutron_map = extract_map_from_logs(caplog)
         assert neutron_map == {
-            0: (32, 33),  # Conv2DStandardV2 (Conv + Relu)
-            1: (35, 36),  # Conv2DDepthwiseV1 (DepthwiseConv + Relu)
-            2: (38, 39),  # Conv2DPointwise (Conv + Relu)
-            3: (41, 42),  # Conv2DDepthwiseV1 (DepthwiseConv + Relu)
-            4: (44, 45),  # Conv2DPointwise (Conv + Relu)
-            5: (47, 48),  # Conv2DDepthwiseV1 (DepthwiseConv + Relu)
-            6: (50, 51),  # Conv2DPointwise (Conv + Relu)
-            7: (53, 54),  # Conv2DDepthwiseV1 (DepthwiseConv + Relu)
-            8: (56, 57),  # Conv2DPointwise (Conv + Relu)
-            9: (59, 60),  # Conv2DDepthwiseV1 (DepthwiseConv + Relu)
-            10: (62, 63),  # Conv2DPointwise (Conv + Relu)
-            11: (65, 66),  # Conv2DDepthwiseV1 (DepthwiseConv + Relu)
-            12: (68, 69),  # Conv2DPointwise (Conv + Relu)
-            13: (71, 72),  # Conv2DDepthwiseV1 (DepthwiseConv + Relu)
-            14: (74, 75),  # Conv2DPointwise (Conv + Relu)
-            15: (77, 78),  # Conv2DDepthwiseV1 (DepthwiseConv + Relu)
-            16: (80, 81),  # Conv2DPointwise (Conv + Relu)
-            17: (83, 84),  # Conv2DDepthwiseV1 (DepthwiseConv + Relu)
-            18: (86, 87),  # Conv2DPointwise (Conv + Relu)
-            19: (89, 90),  # Conv2DDepthwiseV1 (DepthwiseConv + Relu)
-            20: (92, 93),  # Conv2DPointwise (Conv + Relu)
-            21: (95, 96),  # Conv2DDepthwiseV1 (DepthwiseConv + Relu)
-            22: (98, 99),  # Conv2DPointwise (Conv + Relu)
-            23: (101, 102),  # Conv2DDepthwiseDense (DepthwiseConv + Relu)
-            24: (104, 105),  # Conv2DPointwise (Conv + Relu)
-            25: (107, 108),  # Conv2DDepthwiseDense (DepthwiseConv + Relu)
-            26: (110, 111),  # Conv2DPointwise (Conv + Relu)
-            27: (112, 113),  # GlobalAvgPool (Mean + Reshape)
-            28: (114,),  # FullyConnected
+            0: (58, 59, 60),  # Conv2DStandardV2 (Pad + Conv + Relu)
+            1: (61, 62),  # Conv2DDepthwiseV1 (DepthwiseConv + Relu)
+            2: (63, 64),  # Conv2DPointwise (Conv + Relu)
+            3: (65, 66, 67),  # Conv2DDepthwiseV1 (Pad + DepthwiseConv + Relu)
+            4: (68, 69),  # Conv2DPointwise (Conv + Relu)
+            5: (70, 71),  # Conv2DDepthwiseV1 (DepthwiseConv + Relu)
+            6: (72, 73),  # Conv2DPointwise (Conv + Relu)
+            7: (74, 75, 76),  # Conv2DDepthwiseV1 (Pad + DepthwiseConv + Relu)
+            8: (77, 78),  # Conv2DPointwise (Conv + Relu)
+            9: (79, 80),  # Conv2DDepthwiseV1 (DepthwiseConv + Relu)
+            10: (81, 82),  # Conv2DPointwise (Conv + Relu)
+            11: (83, 84, 85),  # Conv2DDepthwiseV1 (Pad + DepthwiseConv + Relu)
+            12: (86, 87),  # Conv2DPointwise (Conv + Relu)
+            13: (88, 89),  # Conv2DDepthwiseV1 (DepthwiseConv + Relu)
+            14: (90, 91),  # Conv2DPointwise (Conv + Relu)
+            15: (92, 93),  # Conv2DDepthwiseV1 (DepthwiseConv + Relu)
+            16: (94, 95),  # Conv2DPointwise (Conv + Relu)
+            17: (96, 97),  # Conv2DDepthwiseV1 (DepthwiseConv + Relu)
+            18: (98, 99),  # Conv2DPointwise (Conv + Relu)
+            19: (100, 101),  # Conv2DDepthwiseV1 (DepthwiseConv + Relu)
+            20: (102, 103),  # Conv2DPointwise (Conv + Relu)
+            21: (104, 105),  # Conv2DDepthwiseV1 (DepthwiseConv + Relu)
+            22: (106, 107),  # Conv2DPointwise (Conv + Relu)
+            23: (108, 109, 110),  # Conv2DDepthwiseDense (Pad + DepthwiseConv + Relu)
+            24: (111, 112),  # Conv2DPointwise (Conv + Relu)
+            25: (113, 114),  # Conv2DDepthwiseDense (DepthwiseConv + Relu)
+            26: (115, 116),  # Conv2DPointwise (Conv + Relu)
+            27: (117, 118),  # GlobalAvgPool (Mean + Reshape)
+            28: (119,),  # FullyConnected
             29: (),  # Neutron Dump
         }
 

backends/nxp/tests/models/test_mlperf_tiny_image_classification.py

@@ -80,14 +80,21 @@ def test_mlperf_tiny_classification_mse_cpu_vs_npu(
         else None
     )
 
+    # Portable constant_pad_nd does not support channels-last tensors.
+    ref_model = (
+        ReferenceModel.QUANTIZED_EDGE_PYTHON
+        if channels_last
+        else ReferenceModel.QUANTIZED_EXECUTORCH_CPP
+    )
+
     lower_run_compare(
         model,
         [input_spec],
         model_verifier,
         request,
         dataset_creator=dataset_creator,
         output_comparator=comparator,
-        reference_model=ReferenceModel.QUANTIZED_EXECUTORCH_CPP,
+        reference_model=ref_model,
         mocker=mocker,
         use_qat=use_qat,
         train_fn=train_fn,

backends/nxp/tests/models/test_mlperf_tiny_keyword_spotting.py

@@ -32,11 +32,11 @@
 BOUNDS_MSE = {
     "PTQ": {
         "channels-last": np.inf,
-        "channels-first": 5.5e-7,
+        "channels-first": 2.0e-6,
     },
     "QAT": {
         "channels-last": np.inf,
-        "channels-first": 3.3e-5,
+        "channels-first": 5.0e-5,
     },
 }
 

backends/nxp/tests/models/test_mlperf_tiny_visual_wake_words.py

@@ -23,14 +23,15 @@
 from executorch.backends.nxp.tests.nsys_testing import (
     lower_run_compare,
     lower_run_compare_ptq_qat,
+    ReferenceModel,
 )
 from executorch.backends.nxp.tests.use_qat import *  # noqa F403
 from executorch.examples.nxp.models.mlperf_tiny.visual_wake_words.mlperf_tiny_visual_wake_words import (
     MLPerfTinyVisualWakeWords,
 )
 
 BOUNDS_MSE = {
-    "PTQ": {"channels-last": 4.6e-8, "channels-first": 5.0e-8},
+    "PTQ": {"channels-last": 1.1e-7, "channels-first": 5.0e-8},
     "QAT": {"channels-last": 3.0e-6, "channels-first": 3.7e-6},
 }
 
@@ -75,13 +76,21 @@ def test_mlperf_tiny_vww_mse_cpu_vs_npu(mocker, request, channels_last, use_qat)
         else None
     )
 
+    # Portable constant_pad_nd does not support channels-last tensors.
+    ref_model = (
+        ReferenceModel.QUANTIZED_EDGE_PYTHON
+        if channels_last
+        else ReferenceModel.QUANTIZED_EXECUTORCH_CPP
+    )
+
     lower_run_compare(
         model,
         [input_spec],
         model_verifier,
         request,
         dataset_creator=dataset_creator,
         output_comparator=comparator,
+        reference_model=ref_model,
         mocker=mocker,
         use_qat=use_qat,
         train_fn=train_fn,