Run the code in single GPU

[shhh0620]

Dear Song,

Hi, thanks for your great work.
I try to reproduce your work to enhance it a bit, but there is some problem in my setting.
I have to use a single GPU due to the limitation of resources. For this, I add
os.environ["CUDA_VISIBLE_DEVICES"]='1'
in main.py or set device with
torch.device('cuda:1')
in run_lib.train.

However, it always assigns gpu:0, so
RuntimeError: cuDNN error: CUDNN_STATUS_MAPPING_ERROR
is occurred.

Is there any solution for it?

GPU that I used is NVIDIA RTX3090 with 24GB.
Thank you.

[arpanpoudel]

You can create a .sh file to export export CUDA_VISIBLE_DEVICES=1 and use config.device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu') on config file. this solved the issue

[shhh0620]

Thank you!