Hi Alban,
Thank for the reply. I don’t think that would work for me as the constraints are such that the forward propagation involves p[0] rather than the entire p matrix and hence p is not directly related in the forward propagation. Do you think there is a way out?