pickle -> parquet

Coda-Research-Group · Jan 31, 2025 · adde797 · adde797
1 parent e289c69
commit adde797
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/training/tests/test_create_embedding.py b/training/tests/test_create_embedding.py
@@ -40,7 +40,7 @@ def test_create_embedding():
         assert os.path.exists(output_path)
         assert os.path.getsize(output_path) > 0
         # load embedding.pkl and check if it has the correct shape
-        df = pd.read_pickle(output_path)
+        df = pd.read_parquet(output_path)
         assert df.shape[0] == len(os.listdir(cif_path))
         assert df.shape[1] == expected_dimensionality