Adjusting test for opset 18

Signed-off-by: Boris Fomitchev <[email protected]>
NVIDIA · Aug 28, 2023 · a8d9e43 · a8d9e43
1 parent 1d01e5c
commit a8d9e43
Showing 1 changed file with 5 additions and 3 deletions.
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)