torch.multiprocessing would result in launching separate processes for each set of sub-neurons (neurons 1:3, 4:6, … in layer B). Although it is possible that these processes get executed on the GPU. But, they would nevertheless be treated as separate processes by the GPU scheduler. I think a more native way to do things would be to launch a single CUDA kernel (with large number of threads) that processes the set of sub-neurons (neurons 1:3, 4:6, … in layer B) in separate threads but as a single process. Also, in this case I can avoid having to do any concatenation, myself.
I have not looked into the speeds for two approaches but I think the approach using torch.multiprocessing might be less efficient (speed-wise). Please correct me if I am wrong.
I agree that for the CUDA kernel approach to work, my partial connectivity needs to have some structure. it can not be absolutely random. And, this is true for my case.