Hi,
I am trying to concatenate the activation from n number of clients before passing to my host-DNN model. How can I build the shuffling layer after the concatenation layer at my host-DNN model so that it can shuffle the activation coming from the clients?
Also, while getting the gradients back from host-DNN, gradients should be unshuffled and passed back.

pytorch_discussion1018×164 25 KB

Should we also implement a backward function for the shuffle layer or it will be taken care by autograd?