Arguments are located on different GPUs; couldn't solve by register_buffer

this error caused by register_buffer

once you access the value in it, there is no way to move it to other device even if you call .to(device)

doesn’t work