Pass grad_output of the last layer to previous layers

I am trying to implement Direct Feedback Alignment with a custom autograd function, a diagram of the algorithm can be found here in subplot c. Basically what I want to do (for example in a Linear layer) is changing the gradient computation in backward function from

grad_input =
grad_weight = grad_output.t().mm(input)


grad_input = grad_output  # Directly pass down grad_output without modifiying
grad_weight = 

But this won’t work as the grad_input does not have the corresponding size to the input.

I was wondering if there is a way to implement this with a custom autograd function. Any hint would be appreciated!