search

yang-song / score_sde_pytorch

PyTorch implementation for Score-Based Generative Modeling through Stochastic Differential Equations (ICLR 2021, Oral)
https://arxiv.org/abs/2011.13456
Apache License 2.0
1.58k stars 295 forks source link

Run the code in single GPU #27

Closed shhh0620 closed 11 months ago

shhh0620 commented 1 year ago

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 commented 11 months ago

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 commented 11 months ago

Thank you!