I have one model that I want to train on multi-gpus. It works well on 0.3.0. But when I changed to 0.4.0. It has the error message as " arguments are located on different GPUs". More strange thing was that the first iteration worked fine(the code passed all backward(), optimizer.step() sentences) and the error happened when entering the backward the second time.
The code is quite complicated and also cannot be published yet. Anyone has some idea how this error happens based on the limited information?
Thanks in advance