small fix

YerevaNN · May 22, 2024 · 32cc5aa · 32cc5aa
1 parent b98a2cb
commit 32cc5aa
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/chemlactica/mol_opt/optimization.py b/chemlactica/mol_opt/optimization.py
@@ -106,9 +106,10 @@ def optimize(
                         **data,
                         **config["generation_config"]
                     )
+                    output_texts.extend(tokenizer.batch_decode(output))
+                    del output
                     gc.collect()
                     torch.cuda.empty_cache()
-                    output_texts.extend(tokenizer.batch_decode(output))
 
                 current_mol_entries = []
                 current_optim_entries = []

diff --git a/chemlactica/mol_opt/utils.py b/chemlactica/mol_opt/utils.py
@@ -96,7 +96,7 @@ def __init__(self, size, validation_perc: float):
     #         self.molecule_entries.pop(rand_ind)
     #     print(f"Dump {num} random elements from pool, num pool mols {len(self)}")
 
-    def add(self, entries: List, diversity_score=1.r9):
+    def add(self, entries: List, diversity_score=1.0):
         assert type(entries) == list
         self.optim_entries.extend(entries)
         self.optim_entries.sort(key=lambda x: x.last_entry, reverse=True)