Out of CUDA memory while tracing a model with jit.trace()

jit
1

I got “out of memory” when I tried to trace a model with jit.
What is interesting is that when I run the model in the test mode it works.

What is the difference between testing and tracing?

2

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.

3

With respect to “test mode”, I meant what you said. And I feel like there seems to be a memory leak in my code then.

Thanks for your reply.