Averaging embedding in training + MeanBackwards Autograd

I’m implementing a neural model for a task where part of the input is supposed to be several categorical variables. At inference time, I wish to embed each attribute separately and then take the mean.

If I average the embeddings at train time similarly as during inference, will a model be able to learn each embedding correctly? I.e. will the model be able to figure out which embedding to update and how even if only a mean of them was fed as input? My intuition tells me that MeanBackwards autograd is supposed to help with this but I’m not sure I understand exactly how that is going to work!

You want to pass multiple inputs in parallel and calculate mean of outputs or you want to calculate mean of inputs then pass it to model?