register_buffer solves my problem. The segmentation default actually comes from other parts. It seems when training on multiple GPUs, we cannot call .cuda() during the forward path, so everything should be registered in the buffer.
Thanks so much for your help!