diff --git a/torch_xla/csrc/xla_graph_executor.cpp b/torch_xla/csrc/xla_graph_executor.cpp index 42e38b92184..5b430810816 100644 --- a/torch_xla/csrc/xla_graph_executor.cpp +++ b/torch_xla/csrc/xla_graph_executor.cpp @@ -122,6 +122,11 @@ XLAGraphExecutor::ComputationCache* CreateComputationCache() { } // namespace +XLAGraphExecutor::~XLAGraphExecutor() { + computation_cache_->Clear(); + delete computation_cache_; +} + auto XLAGraphExecutor::DeviceContextArena::Get() -> DeviceContextArena* { static DeviceContextArena* arena = new DeviceContextArena(); return arena; diff --git a/torch_xla/csrc/xla_graph_executor.h b/torch_xla/csrc/xla_graph_executor.h index 5e1acaa9d47..dde81aa2b74 100644 --- a/torch_xla/csrc/xla_graph_executor.h +++ b/torch_xla/csrc/xla_graph_executor.h @@ -30,6 +30,8 @@ namespace torch_xla { class XLAGraphExecutor : public torch::lazy::LazyGraphExecutor { public: + ~XLAGraphExecutor(); + static XLAGraphExecutor* Get(); // Override to use our own DeviceContextArena.