Adjusting test for opset 18

borisfom · borisfom · commit a8d9e439e7d1 · 2023-08-28T15:45:30.000-07:00
Signed-off-by: Boris Fomitchev &lt;bfomitchev@nvidia.com&gt;
diff --git a/tests/L0/run_fused_layer_norm/test_fused_layer_norm.py b/tests/L0/run_fused_layer_norm/test_fused_layer_norm.py
@@ -268,9 +268,11 @@ def _verify_export(self, fused, fused_x):
         # check that export() is working
         onnx_str = torch.onnx.export_to_pretty_string(fused, (fused_x,),
                                                       input_names=['x_in'],
+                                                      opset_version=18,
         )
+        print(onnx_str)
         assert 'x_in' in onnx_str
-        assert 'ReduceMean' in onnx_str
+        assert 'ReduceMean' in onnx_str or 'LayerNormalization' in onnx_str
 
     def test_rms_export(self):
         batch_size = 16
@@ -279,7 +281,7 @@ def test_rms_export(self):
             normalized_shape=normalized_shape, elementwise_affine=True
         ).cuda()
         fused_m = MixedFusedRMSNorm(
-            normalized_shape=normalized_shape, elementwise_affine=True
+            normalized_shape=normalized_shape
         ).cuda()
         native_x, fused_x = _prep_inputs(batch_size, normalized_shape, torch.float32)
         self._verify_export(fused, fused_x)
@@ -292,7 +294,7 @@ def test_layer_norm_export(self):
             normalized_shape=normalized_shape, elementwise_affine=True
         ).cuda()
         fused_m = MixedFusedLayerNorm(
-            normalized_shape=normalized_shape, elementwise_affine=True
+            normalized_shape=normalized_shape
         ).cuda()
         native_x, fused_x = _prep_inputs(batch_size, normalized_shape, torch.float32)
         self._verify_export(fused, fused_x)
