feat: call torch.cuda.empty_cache() per design when empty_cache_per_design=True

leonardomarino · leonardomarino · commit e8bd8d143575 · 2026-04-21T10:59:41.000-04:00
diff --git 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")
 
 

Original file line number	Diff line number	Diff line change
`@@ -188,6 +188,9 @@ def main(conf: HydraConfig) -> None:`
`188`	`188`	`chain_ids=sampler.chain_idx,`
`189`	`189`	`)`
`190`	`190`
	`191`	`+ if conf.inference.empty_cache_per_design and torch.cuda.is_available():`
	`192`	`+ torch.cuda.empty_cache()`
	`193`	`+`
`191`	`194`	`log.info(f"Finished design in {(time.time()-start_time)/60:.2f} minutes")`
`192`	`195`
`193`	`196`