Doesn't work in torch.inference_mode()?

[LyuJZ]

❓ Questions and Help

Hi, thanks a lot this amazing work.

I am integrating the theseus layer to actor-critic framework. However, the theseus layer seems doesn't work under the with torch.inference_mode(): condition. It can only work under the with torch.no_grad(): condition.

Could you give me any explanations about this?

[luisenp]

Hi @LyuJZ. I apologize, but we never tested with torch.inference_mode(). Could you share more details of the error you are getting?

[LyuJZ]

Hi @luisenp , that's all right!
In the torch.inference_mode(), there are no errors. However, from the output of layer.forward(input_dict, optimizer_kwargs={"track_best_solution": True, "verbose": True, "ellipsoidal_damping": True}), it says there are only two iteration steps, and the total loss doesn't decrease at tall.

If I use with torch.no_grad():, the log information says that there are several iteration steps, and the total loss is decreasing consistently.

[luisenp]

Sorry for the delay, somehow I never got notified about this response. Unfortunately, I'm not too familiar with torch.inference_mode() so I don't have a good idea of what can be causing this. Is this issue blocking your progress when using Theseus?

[Ribosome-rbx]

@LyuJZ I'm facing the same issue, have you fixed it?

[Ribosome-rbx]

Some updates on how I solved my problem.
I want to run theseus in the validation/test loop of pytorch_lightning, where the whole pipeline is under the context of "torch.inference_mode()". There will not be any error message from theseus side, and the optimization run only one step and terminate, without any loss decreasing (probably because of the missing grad info). To solve this, simply adding the decorator torch.inference_mode(False) to the optimization function works for me.