search

sniklaus / pytorch-pwc

a reimplementation of PWC-Net in PyTorch that matches the official Caffe version
GNU General Public License v3.0
608 stars 122 forks source link

Cannot select the GPU index #22

Closed v-iashin closed 5 years ago

v-iashin commented 5 years ago

Hi.

It seems that I cannot make the code to run on another GPU rather than 0-th. The following line doesn't make any difference in my case https://github.com/sniklaus/pytorch-pwc/blob/f6138900578214ab4e3daef6743b88f7824293be/run.py#L25

I managed to solve it by introducing the device variable instead and adding torch.cuda.set_device(device) right after the FunctionCorrelation signature https://github.com/sniklaus/pytorch-pwc/blob/f6138900578214ab4e3daef6743b88f7824293be/correlation/correlation.py#L392-L393 and also passing this variable as an argument here https://github.com/sniklaus/pytorch-pwc/blob/f6138900578214ab4e3daef6743b88f7824293be/run.py#L196 and here https://github.com/sniklaus/pytorch-pwc/blob/f6138900578214ab4e3daef6743b88f7824293be/run.py#L204

If this is a bug do you want me to form a Pull Request?

Anyway, thanks a lot for sharing the code. It is really helpful when the official version is on CUDA 8 and PyTorch 0.2!

sniklaus commented 5 years ago

You are correct, it should be set_device instead of just device. Thank you for letting me know! No need to issue a pull request, I just updated it and removed that line altogether. It is usually better to set CUDA_VISIBLE_DEVICES and I only had this line to avoid issues with multi-GPU systems.