When you say “test mode”, do you mean calling eval() and then forward()? There are minor differences; certain Modules have different behavior in eval vs. training mode. In addition, if during testing you aren’t recording gradients (e.g. you set requires_grad to off), you don’t have to build up an autograd graph and keep the bookkeeping structures around.
Keep in mind that tracing will actually execute your model in order to trace it. When you run the model forward on the example inputs manually, does it still oom?
It would be helpful to have a script that we can use to reproduce the issue, otherwise there isn’t a ton of information to go off of.