jac_loss has no gradient for training.

[QuietNoviceCoder]

this is my code :
self.deq = get_deq(
ift = True,
f_solver=deq_solver,
b_solver=deq_solver,
f_max_iter=deq_f_max_iter,
b_max_iter=deq_b_max_iter,
f_tol=deq_f_tol,
b_tol=deq_b_tol,
stop_mode = 'rel',
f_anderson_m=deq_anderson_m if deq_solver == 'anderson' else None,
b_anderson_m=deq_anderson_m if deq_solver == 'anderson' else None,
)

solve , deq_info = self.deq(f,y)

then:
new_z_star = f(z_star.requires_grad_())
jac_loss = jac_reg(new_z_star,z_star)

so that i add the jac_loss to my loss,but it not works.

[sablin39]

Did you add the jac_loss to your overall loss before backward like this? Also, this is a double-backward operation, so it cannot work with torch.compile, and potential bugs may exist.

[QuietNoviceCoder]

Did you add the jac_loss to your overall loss before backward like this? Also, this is a double-backward operation, so it cannot work with torch.compile, and potential bugs may exist.

Yes, I did add jac_loss to the overall loss before calling backward, but I couldn't observe the gradients.

However, if I use only jac_loss, the internal parameters do update. This implies that the gradients exist and are working, which is exactly what I find so strange.

[sablin39]

Did you add the jac_loss to your overall loss before backward like this? Also, this is a double-backward operation, so it cannot work with torch.compile, and potential bugs may exist.

Yes, I did add jac_loss to the overall loss before calling backward, but I couldn't observe the gradients.

However, if I use only jac_loss, the internal parameters do update. This implies that the gradients exist and are working, which is exactly what I find so strange.

I've tried to replicate your situation with one cross attn layer wrapped by DEQ block, but I still get the jac_loss (though very small at ~2e-6, which is common for randomly initialized models). Please provide a more detailed training setting (e.g., overall loss calculation, network structure), or switch from IFT to phantom gradient and see if the error persists. You may also try visualizing your computation graph with torchviz.