pub_soft/RFdiffusion
3
0
Fork 0
mirror of https://github.com/RosettaCommons/RFdiffusion.git synced 2026-06-04 18:44:21 +08:00
Code Issues Packages Projects Releases Wiki Activity

Add verbose checking for GPU

Browse Source
This commit is contained in:
Nathaniel Bennett
2023-04-06 13:58:48 -07:00
committed by Joseph Watson/Watchwell
parent 6adcf18465
commit 92b83decf3
1 changed files with 9 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
scripts/run_inference.py
View File

@@ -41,6 +41,15 @@ def main(conf: HydraConfig) -> None:
if conf.inference.deterministic:
make_deterministic()
# Check for available GPU and print result of check
if torch.cuda.is_available():
device_name = torch.cuda.get_device_name(torch.cuda.current_device())
log.info(f"Found GPU with device_name {device_name}. Will run RFdiffusion on {device_name}")
else:
log.info("////////////////////////////////////////////////")
log.info("///// NO GPU DETECTED! Falling back to CPU /////")
log.info("////////////////////////////////////////////////")
# Initialize sampler and target/contig.
sampler = iu.sampler_selector(conf)