I’m currently implementing a network architecture where I have to pass my data through a certain layer for multiple times.
The network can be simplified as:
Input ->  Layer A -> Layer B -> Layer A -> Output

However I dont want Layer A to be updated by the gradient, i.e. freezed, on the latter call, while still be updated by the first call. Is this possible with pytorch? I’ve tried to come up with something using no_grad, requires_grad or detach() but I couldn’t.
A hacky workaround Im thinking of is to do gradient scaling during the backward pass:
Input Gradients <- Layer A <- Layer B <- Scale up gradient by λ <- Layer A <- Scale down gradient by λ <- Output Gradients

Is there a cleaner implementation? Thanks.