AFAICT this is exactly tthe case where tracing doesn’t work well. See the Warning section in the torch.jit.trace documentation page:
https://pytorch.org/docs/stable/generated/torch.jit.trace.html#torch.jit.trace 3

…
Tracing will not record any control-flow like if-statements or loops.
…

It also offers a possible solution:

In cases like these, tracing would not be appropriate and scripting 2 is a better choice. If you trace such models, you may silently get incorrect results on subsequent invocations of the model. The tracer will try to emit warnings when doing something that may cause an incorrect trace to be produced.