From 0966015bc74c4fc775c055840780d18ce19cbef5 Mon Sep 17 00:00:00 2001
From: Priya Kasimbeg <kasimbeg@google.com>
Date: Sat, 18 Nov 2023 03:54:47 +0000
Subject: [PATCH] debug

---
 tests/modeldiffs/criteo1tb_resnet/compare.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tests/modeldiffs/criteo1tb_resnet/compare.py b/tests/modeldiffs/criteo1tb_resnet/compare.py
index b5b0cfb62..0253f65bc 100644
--- a/tests/modeldiffs/criteo1tb_resnet/compare.py
+++ b/tests/modeldiffs/criteo1tb_resnet/compare.py
@@ -18,8 +18,8 @@
 def key_transform(k):
   print('key transform: ')
   new_key = []
-  s_count = 0
-  resnet_count = 0
+  s_count = None
+  resnet_count = None
   print(k)
   for i in k:
     print(f'in transform {i}')
@@ -30,7 +30,7 @@ def key_transform(k):
       return ('embedding_table',)
     if 'ResNetBlock' in i:
       name, count = i.split('_')
-      resnet_count = resnet_count + 1 
+      resnet_count = int(count)
       continue
     if 'Linear' in i:
       i = i.replace('Linear', 'Dense')