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.