OK, I felt a little frustrated. I did everything I could to reduce GPU memory. But the usage of my pytorch code is still more than twice memory-consuming than my tensorflow implementation. Because the code is confidential, I cannot release it in public.
My network B is an inception v3 network. I found that after I passed the output of A to B, the GPU memory increases 4 GB. The batch size is only about 20.