diff --git a/config/inference/base.yaml b/config/inference/base.yaml index 3bb0a5c1..686e71bf 100644 --- a/config/inference/base.yaml +++ b/config/inference/base.yaml @@ -11,6 +11,7 @@ inference: model_only_neighbors: False output_prefix: samples/design write_trajectory: True + empty_cache_per_design: False scaffold_guided: False model_runner: SelfConditioning cautious: True diff --git a/scripts/run_inference.py b/scripts/run_inference.py index 3fb6466e..3ebb5e3f 100755 --- a/scripts/run_inference.py +++ b/scripts/run_inference.py @@ -188,6 +188,9 @@ def main(conf: HydraConfig) -> None: chain_ids=sampler.chain_idx, ) + if conf.inference.empty_cache_per_design and torch.cuda.is_available(): + torch.cuda.empty_cache() + log.info(f"Finished design in {(time.time()-start_time)/60:.2f} minutes")