sort weights for sparse

VowpalWabbit · Oct 5, 2023 · f22aff2 · f22aff2
1 parent eb3ee41
commit f22aff2
Showing 1 changed file with 8 additions and 5 deletions.
diff --git a/test/run_tests_model_gen_and_load.py b/test/run_tests_model_gen_and_load.py
@@ -179,12 +179,15 @@ def load_model(
         weights_dir = working_dir / "test_weights"
         weights_dir.mkdir(parents=True, exist_ok=True)
         weight_file = str(weights_dir / f"weights_{test_id}.json")
+
+        # Load old weights
         old_weights = json.load(open(weight_file))
-        if new_weights != old_weights:
-            print("OLD WEIGHTS:")
-            print(new_weights)
-            print("NEW WEIGHTS:")
-            print(old_weights)
+
+        # Sort the weights inside both dictionaries based on 'index'
+        new_weights['weights'].sort(key=lambda x: x['index'])
+        old_weights['weights'].sort(key=lambda x: x['index'])
+
+        # Assert if both sorted weights are equal
         assert new_weights == old_weights
         vw.finish()
     except Exception as e: