Delete the ref_pos entry only after the pickle is written

PiperOrigin-RevId: 946977350
Change-Id: I033f1e72d88d0a54713838baaecb1ac41e398031

Author: Augustin-Zidek

run_alphafold_data_test.py

@@ -226,10 +226,11 @@ def test_featurisation(self):
         ccd=chemical_components.Ccd(),
         buckets=None,
     )
-    del featurised_example[0]['ref_pos']  # Depends on specific RDKit version.
 
     with _output('featurised_example.pkl') as (_, output):
       output.write(pickle.dumps(featurised_example))
+
+    del featurised_example[0]['ref_pos']  # Depends on specific RDKit version.
     featurised_example = jax.tree_util.tree_map(_hash_data, featurised_example)
     with _output('featurised_example.json') as (result_path, output):
       output.write(