From 6913e9225989a97d54d105087cd07742e25901e5 Mon Sep 17 00:00:00 2001 From: Priya Kasimbeg Date: Wed, 22 Nov 2023 03:15:29 +0000 Subject: [PATCH] fix --- tests/modeldiffs/criteo1tb_layernorm/compare.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/modeldiffs/criteo1tb_layernorm/compare.py b/tests/modeldiffs/criteo1tb_layernorm/compare.py index fc306e625..130d0d84b 100644 --- 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')