Skip to content

Commit a9a1355

Browse files
authored
[FxImporter] remove weakref finalizer of reftracker (#3995)
Fixes iree-org/iree-turbine#281. **_TL;DR:_** The `weakref.finalize` objects cause the model parameters to be kept in memory in-between consecutive `aot.export` calls in the same process. We remove them to enable releasing the memory, this does not change the behavior of `RefTracker` or `RefMapping` classes anyhow else. Signed-off-by: Ege Beysel <[email protected]>
1 parent 217d0a8 commit a9a1355

File tree

1 file changed

+6
-4
lines changed

1 file changed

+6
-4
lines changed

python/torch_mlir/extras/fx_importer.py

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2247,13 +2247,15 @@ def track(self, referrent: Any) -> RefMapping:
22472247
if existing:
22482248
return existing
22492249
info = RefMapping(referrent)
2250-
if referrent is not Empty:
2251-
weakref.finalize(referrent, self._ref_finalizer, ref_id)
2250+
# Finalizer is removed due to a memory leak
2251+
# See: https://github.com/iree-org/iree-turbine/issues/281
2252+
# if referrent is not Empty:
2253+
# weakref.finalize(referrent, self._ref_finalizer, ref_id)
22522254
self._refs[ref_id] = info
22532255
return info
22542256

2255-
def _ref_finalizer(self, ref_id: int):
2256-
del self._refs[ref_id]
2257+
# def _ref_finalizer(self, ref_id: int):
2258+
# del self._refs[ref_id]
22572259

22582260

22592261
################################################################################

0 commit comments

Comments
 (0)