fix

mlcommons · Nov 22, 2023 · 6913e92 · 6913e92
1 parent c44b5bd
commit 6913e92
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/tests/modeldiffs/criteo1tb_layernorm/compare.py b/tests/modeldiffs/criteo1tb_layernorm/compare.py
@@ -30,6 +30,9 @@ def key_transform(k):
       i = i.replace('Linear', 'Dense')
       name, count = i.split('_')
       i = name + '_' + str(s_count * 3 + int(count))
+    if 'LayerNorm' in i:
+      name, count = i.split('_')
+      i = name + '_' + str(s_count * 3 + int(count))
     elif 'weight' in i:
       i = i.replace('weight', 'kernel')