Guidance Needed: Specifying GPU for AF2_initial_guess

[assiduous006]

I am currently working with the AF2_initial_guess application on a Ubuntu machine equipped with multiple GPUs. I am interested in designating a specific GPU for running this application to optimize performance.

Could anyone provide detailed instructions or point me towards relevant resources that could assist me in achieving this setup? Your help would be greatly appreciated.

Thank you in advance for your time and assistance.

[nrbennet]

In the code there are references to device. Currently this is just torch.device('cuda') so it picks a GPU but not a specific one. If you want to specify which gpu to use you can choose a rank of GPU, so if you want the first GPU according to CUDA you can do this torch.device('cuda:0'). You will just need to code up something to correctly choose an unused rank during runtime

[assiduous006]

Thanks for your immediate response!

I am currently working with the AF2_initial_guess application, specifically running “predict.py”. The application appears to use TensorFlow and JAX as its backend.

However, I have been unable to find any options within the code that would allow me to designate a specific GPU.

Could you please provide guidance on how this could be achieved, or if there are any workarounds available?

Any assistance would be greatly appreciated. Thank you in advance for your time and help.

[nrbennet]

A bit hacky but this should work:

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"  # Use the first GPU

[assiduous006]

I’m eager to put this to the test! Your support is invaluable, and I thank you sincerely for your guidance.