Merge branch 'master' into tagdataset/add-llm-exp-pred

pyg-team · Jan 10, 2025 · ac7bfed · ac7bfed
2 parents ac402f3 + ef02854
commit ac7bfed
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/torch_geometric/datasets/molecule_gpt_dataset.py b/torch_geometric/datasets/molecule_gpt_dataset.py
@@ -371,6 +371,7 @@ def extract_one_SDF_file(block_id: int) -> None:
                     writer.write(mol)
                     valid_mol_count += 1
 
+                writer.close()
                 print(f"block id: {block_id}\nfound {valid_mol_count}\n\n")
                 sys.stdout.flush()
                 return
@@ -410,6 +411,7 @@ def extract_one_SDF_file(block_id: int) -> None:
                 print(f"block id: {block_id} with 0 valid SDF file")
                 continue
 
+        writer.close()
         print(f"In total: {len(found_CID_set)} molecules")
 
         # Step 05. Convert to PyG data format