debugging

mlcommons · Dec 7, 2023 · 4ba6736 · 4ba6736
1 parent 59dfe11
commit 4ba6736
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/tests/modeldiffs/torch2jax_utils.py b/tests/modeldiffs/torch2jax_utils.py
@@ -77,6 +77,7 @@ def key_transform(self, k_transform_fn):
     }
 
   def value_transform(self, v_transform_fn):
+    print(self.flattened_jax_model.keys())
     self.pytorch_sd = {
         k: v_transform_fn(k, self.pytorch_sd[k], self.flattened_jax_model[k])
         for k in self.pytorch_sd