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?