bayesianflow-for-chem 2.2.2__tar.gz → 2.2.4__tar.gz

{bayesianflow_for_chem-2.2.2 → bayesianflow_for_chem-2.2.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bayesianflow_for_chem
-Version: 2.2.2
+Version: 2.2.4
 Summary: Bayesian flow network framework for Chemistry
 Home-page: https://augus1999.github.io/bayesian-flow-network-for-chemistry/
 Author: Nianze A. Tao

{bayesianflow_for_chem-2.2.2 → bayesianflow_for_chem-2.2.4}/bayesianflow_for_chem/__init__.py

@@ -16,7 +16,7 @@ __all__ = [
     "MLP",
     "EnsembleChemBFN",
 ]
-__version__ = "2.2.2"
+__version__ = "2.2.4"
 __author__ = "Nianze A. Tao (Omozawa Sueno)"
 
 

{bayesianflow_for_chem-2.2.2 → bayesianflow_for_chem-2.2.4}/bayesianflow_for_chem/cli.py

@@ -569,7 +569,7 @@ def main_script(version: str) -> None:
         model = Model(bfn, mlp, scorer)
         model.model.semi_autoregressive = runtime_config["train"]["semi_autoregressive"]
         # ####### strat training #######
-        os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128"
+        os.environ["PYTORCH_ALLOC_CONF"] = "max_split_size_mb:128"
         if not runtime_config["train"]["dynamic_padding"]:
             os.environ["MAX_PADDING_LENGTH"] = f"{lmax}"  # important!
         torch.set_float32_matmul_precision("medium")
@@ -696,7 +696,7 @@ def main_script(version: str) -> None:
         with open(runtime_config["inference"]["result_file"], "w") as f:
             f.write("\n".join(mols))
     # ------- finished -------
-    print(" ####### job finished #######")
+    print("*" * 25 + " job finished " + "*" * 25)
     print(_END_MESSAGE)
 
 

{bayesianflow_for_chem-2.2.2 → bayesianflow_for_chem-2.2.4}/bayesianflow_for_chem/model.py

@@ -1171,6 +1171,9 @@ class EnsembleChemBFN(ChemBFN):
         assert (
             isinstance(c, dict) is self._label_is_dict
         ), f"`c` should be a {'`dict` instance' if self._label_is_dict else '`list` instance'} but got {type(c)} instand."
+        assert len(c) == len(
+            self.models
+        ), f"Number of conditions should match the number of LoRA models. We have {len(self.models)} LoRA models but {len(c)} conditions were provided."
         out: Dict[str, Tensor] = {}
         if isinstance(c, list):
             c = dict(zip([f"val_{i}" for i in range(len(c))], c))

{bayesianflow_for_chem-2.2.2 → bayesianflow_for_chem-2.2.4}/bayesianflow_for_chem.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bayesianflow_for_chem
-Version: 2.2.2
+Version: 2.2.4
 Summary: Bayesian flow network framework for Chemistry
 Home-page: https://augus1999.github.io/bayesian-flow-network-for-chemistry/
 Author: Nianze A. Tao