Update data_correctness.py