# Torch.jit.trace() memory leak workarounds

**URL:** <https://discuss.pytorch.org/t/torch-jit-trace-memory-leak-workarounds/199184>\
**Category:** jit\
**Created:** [March 19, 2024, 1:44pm UTC](https://discuss.pytorch.org/t/torch-jit-trace-memory-leak-workarounds/199184 "2024-03-19T13:44:44Z")\
**Posts on this page:** 5\
**Page:** 1

<div class="post-metadata">

**Author:** ![dshin](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/dshin/32/54230_2.png) [@dshin](https://discuss.pytorch.org/u/dshin)\
**Post date:** [March 19, 2024, 1:44pm UTC](https://discuss.pytorch.org/t/torch-jit-trace-memory-leak-workarounds/199184/1 "2024-03-19T13:44:44Z")

</div>

It has been reported [elsewhere](https://discuss.pytorch.org/t/possible-memory-leak-in-torch-jit-trace/74455) that `torch.jit.trace()` has a memory leak. This is a problem for me because my application repeatedly calls `torch.jit.trace().save()` in a loop (later loading the saved model with torchlib in c++).

One workaround suggested [here](https://github.com/pytorch/pytorch/issues/35600#issuecomment-894807668) is to wrap the `torch.jit.trace()` call in a separate process by using `futures.ProcessPoolExecutor`.

This appears to work well in single-threaded settings. However, in multi-threaded settings, this occasionally leads to deadlock, due to a variant of this [issue](https://github.com/python/cpython/issues/50970). In general, forking multi-threaded processes in python is known to be quite fragile.

I am looking for a workaround. Is there a good way to `torch.jit.trace().save()` in a loop without deadlocking and without incurring a memory leak in a multi-threaded application?

If the answer is, “upgrade to pytorch 2.0 and use `torch.compile()`”, I’m ok with that. Anything that works.

---

<div class="post-metadata">

**Author:** ![ptrblck](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/32/1823_2.png) [@ptrblck](https://discuss.pytorch.org/u/ptrblck)\
**Post date:** [March 20, 2024, 1:59pm UTC](https://discuss.pytorch.org/t/torch-jit-trace-memory-leak-workarounds/199184/2 "2024-03-20T13:59:27Z")

</div>

> [@dshin](#):
>
> If the answer is, “upgrade to pytorch 2.0 and use `torch.compile()`”, I’m ok with that. Anything that works.

That might indeed be the recommendation here, since `TorchScript` is in maintenance mode and will thus not receive any major updates or big fixes anymore. I’m also hesitant to look for other workarounds as the risk of running into a new issue, that won’t be fixed, will still be there.

---

<div class="post-metadata">

**Author:** ![dshin](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/dshin/32/54230_2.png) [@dshin](https://discuss.pytorch.org/u/dshin)\
**Post date:** [March 20, 2024, 2:12pm UTC](https://discuss.pytorch.org/t/torch-jit-trace-memory-leak-workarounds/199184/3 "2024-03-20T14:12:34Z")

</div>

Thanks for the reply. Guess it’s time to upgrade!

---

<div class="post-metadata">

**Author:** ![dshin](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/dshin/32/54230_2.png) [@dshin](https://discuss.pytorch.org/u/dshin)\
**Post date:** [March 25, 2024, 2:27pm UTC](https://discuss.pytorch.org/t/torch-jit-trace-memory-leak-workarounds/199184/4 "2024-03-25T14:27:44Z")

</div>

I started to look into the torch 2.0 documentation a bit, and am a bit confused. The latest page on exporting models to c++ appears to be this [one](https://pytorch.org/docs/stable/cpp_index.html), which still talks about using TorchScript. But I understood your answer to mean that upgrading to torch 2.0 will make it so that TorchScript won’t be needed anymore to export a model to c++. What gives?

Is the recommended approach to export to c++ to export in ONNX format, now? And if so, on the c++ side, to use a non-torch library like ONNX Runtime?

---

<div class="post-metadata">

**Author:** ![dshin](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/dshin/32/54230_2.png) [@dshin](https://discuss.pytorch.org/u/dshin)\
**Post date:** [April 3, 2024, 3:39pm UTC](https://discuss.pytorch.org/t/torch-jit-trace-memory-leak-workarounds/199184/5 "2024-04-03T15:39:04Z")

</div>

There was in fact a simple fix to my original problem. Instead of doing the `torch.jit.trace()` call in a separate process by using `futures.ProcessPoolExecutor`, do it with `subprocess`, pickling the arguments to disk and then reloading in the subprocess. This gets around the multi-threaded + fork sensitivities.
