Straight Through Estimator like

Hello,
I am trying to implement the Straight through Estimator method mentioned in this paper: https://arxiv.org/pdf/2209.04836.pdf. In simple words, there are two models of the same architecture and I am trying to build a third model that is the average of model1 and the non-differentiable operation of model2 (so I would average corresponding weight matrices). However when I back pass I would like to do it only on model2.

I tried the following:
model3 = type(model1)()
for p1, p2_orig, p2, p3 in zip(model1.parameters(), model2.parameters() non_differentiable_func(model2).parameters(), model3.parameters()):
ps.copy_(p1) = 0.5*(p1.requires_grad_(False) + p2.requires_grad_(False) +p2_orig - p2_orig.requires_grad_(False)

However, this throws an error stating that values can’t be modified in leaf tensors or something like that, but as far as I understood I can’t modify the contents as such.

Any help on this matter would be appreciated