Computers with dual graphics card configurations will report errors when run

[Endygao25]

It has been deployed locally. However, I encountered a technical problem, my computer is installed with dual graphics cards, so it will report an error when running, prompting:
"RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1! "

I tried to change a command line in the file ""memory_management.py"":
gpu = torch.device('cuda') to gpu = torch.device('cuda:0').
But still no effect.

How to remedy this situation? I'm looking forward to hearing from you!

[ayuLiao]

you need change like this.

I don't have time to submit a PR, you can revise it by yourself.