OOM in Python 3.12, but not in 3.10

[lyhyl]

🐛 Bug description

import gc
import torch
import torch.nn as nn
from ignite.engine import create_supervised_trainer, Events
from ignite.handlers import ProgressBar
from torch.optim import Adam
from torch.utils.data import DataLoader, TensorDataset

def do(device, model):
    x = torch.rand(32, 1, 64, 64, 32)
    y = torch.round(torch.rand(32, 1))
    ds = TensorDataset(x, y)
    train_loader = DataLoader(ds, 6)
    optim = Adam(model.parameters(), 1e-4)
    loss = nn.BCEWithLogitsLoss()
    trainer = create_supervised_trainer(model, optim, loss, device)
    pbar = ProgressBar()
    pbar.attach(trainer)
    trainer.run(train_loader, 1)
    pbar.close()

def main():
    device = torch.device("cuda")
    for i in range(100):
        model = nn.Sequential(nn.Flatten(), nn.Linear(64*64*32, 6500), nn.ReLU(), nn.Linear(6500, 1))
        model = model.to(device)
        do(device, model)
        del model
        gc.collect()
        torch.cuda.empty_cache()
        print("!!!", i, torch.cuda.memory_allocated(), torch.cuda.max_memory_allocated())

if __name__ == "__main__":
    main()

This code get OOM in python 3.12, but works fine in python 3.10.
Circular reference seems not being handled properly.

The first run in main loop will be fine, but the resources are not released, then the second run will OOM.

Environment

• PyTorch Version: 2.7.1+cu128
• Ignite Version: 0.5.2
• OS (e.g., Linux): ubuntu 24
• How you installed Ignite: uv
• Python version: 3.12
• Any other relevant information: RTX 4090 24G, NVIDIA-SMI 575.64.03 Driver Version: 575.64.03 CUDA Version: 12.9

[lyhyl]

A workaround, for all trainer (create_supervised_trainer) and evaluator (create_supervised_evaluator), clear their _event_handlers

    trainer._event_handlers.clear()
    evaluator._event_handlers.clear()

[vfdev-5]

Hey @lyhyl thanks for reporting this!
I wonder whether the OOM is there without the progress bar?
When the OOM is happening on which line?

[lyhyl]

Hey @lyhyl thanks for reporting this! I wonder whether the OOM is there without the progress bar? When the OOM is happening on which line?

Anything creates circular reference will causes GPU memory leak (at least not freed on time), not only ProgressBar, but also TensorboardLogger, metrics in evaluator.

Thus, GPU memory usage gets higher and higher util some point that need to allocate new GPU memory (most of case, trainer.run), then OOM happen.

Another case, I make network smaller to demonstrate the progress:

import torch
import torch.nn as nn
from ignite.engine import Events, create_supervised_trainer, create_supervised_evaluator
from ignite.handlers import ProgressBar, TensorboardLogger
from ignite.handlers.tensorboard_logger import OptimizerParamsHandler
from torch.optim import Adam
from ignite.metrics import Loss
from torch.utils.data import DataLoader, TensorDataset

def do(model, dataloader, device):
    optim = Adam(model.parameters(), 1e-4)
    loss = nn.BCEWithLogitsLoss()
    trainer = create_supervised_trainer(model, optim, loss, device)
    metrics = {"Loss": Loss(loss)}
    evaluator = create_supervised_evaluator(model, metrics, device)

    pbar = ProgressBar()
    pbar.attach(trainer)

    tb_logger = TensorboardLogger(log_dir="runs")
    tb_logger.attach(trainer, OptimizerParamsHandler(optim), Events.EPOCH_STARTED)

    trainer.run(dataloader, 1)
    @trainer.on(Events.COMPLETED)
    def completed(engine):
        evaluator.run(dataloader)

    tb_logger.close()
    pbar.close()
    # trainer._event_handlers.clear()
    # evaluator._event_handlers.clear()

def main():
    device = torch.device("cuda")
    x = torch.rand(32, 1, 64, 64, 32)
    y = torch.round(torch.rand(32, 1))
    ds = TensorDataset(x, y)
    dataloader = DataLoader(ds, 6)
    for i in range(100):
        N = 3000
        model = nn.Sequential(nn.Flatten(), nn.Linear(64 * 64 * 32, N), nn.ReLU(), nn.Linear(N, 1))
        model = model.to(device)
        do(model, dataloader, device)
        print("!!!", i, torch.cuda.memory_allocated(), torch.cuda.max_memory_allocated())

if __name__ == "__main__":
    main()

I get following in console:

!!! 0 6309382144 7888638464                                                                                                                                                                                                                        
!!! 1 6309382144 14176974848                                                                                                                                                                                                                       
!!! 2 12600938496 14180194816                                                                                                                                                                                                                      
!!! 3 18892494848 20471751168                                                                                                                                                                                                                      
!!! 4 18892494848 20471751168                                                                                                                                                                                                                      
Current run is terminating due to exception: CUDA out of memory. Tried to allocate 1.46 GiB. GPU 0 has a total capacity of 23.51 GiB of which 224.88 MiB is free. Including non-PyTorch memory, this process has 22.45 GiB memory in use. Of the allocated memory 21.99 GiB is allocated by PyTorch, and 1.66 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)
Engine run is terminating due to exception: CUDA out of memory. Tried to allocate 1.46 GiB. GPU 0 has a total capacity of 23.51 GiB of which 224.88 MiB is free. Including non-PyTorch memory, this process has 22.45 GiB memory in use. Of the allocated memory 21.99 GiB is allocated by PyTorch, and 1.66 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)

and the solution is:

    trainer._event_handlers.clear()
    evaluator._event_handlers.clear()

[goanpeca]

@lyhyl thanks for the details !

It seems it is indeed something that changed in python handling circular references in the event handling system.

Engine._event_handlers[event] = [
    (handler, (self, *args), kwargs)  # ← 'self' creates Engine → Handler → Engine cycle
]

Could you test installing this PR branch #3447

pip install git+https://github.com/gaonpeca/ignite.git@fix/oom and now use the context manager, or the close method?

import torch
import torch.nn as nn
from ignite.engine import Events, create_supervised_trainer, create_supervised_evaluator
from ignite.handlers import ProgressBar, TensorboardLogger
from ignite.handlers.tensorboard_logger import OptimizerParamsHandler
from torch.optim import Adam
from ignite.metrics import Loss
from torch.utils.data import DataLoader, TensorDataset

def do(model, dataloader, device):
    optim = Adam(model.parameters(), 1e-4)
    loss = nn.BCEWithLogitsLoss()
    trainer = create_supervised_trainer(model, optim, loss, device)
    metrics = {"Loss": Loss(loss)}
    evaluator = create_supervised_evaluator(model, metrics, device)

    pbar = ProgressBar()
    pbar.attach(trainer)

    tb_logger = TensorboardLogger(log_dir="runs")
    tb_logger.attach(trainer, OptimizerParamsHandler(optim), Events.EPOCH_STARTED)

    trainer.run(dataloader, 1)
    @trainer.on(Events.COMPLETED)
    def completed(engine):
        evaluator.run(dataloader)

    tb_logger.close()
    pbar.close()
    trainer.close()
    evaluator.close()

def main():
    device = torch.device("cuda")
    x = torch.rand(32, 1, 64, 64, 32)
    y = torch.round(torch.rand(32, 1))
    ds = TensorDataset(x, y)
    dataloader = DataLoader(ds, 6)
    for i in range(100):
        N = 3000
        model = nn.Sequential(nn.Flatten(), nn.Linear(64 * 64 * 32, N), nn.ReLU(), nn.Linear(N, 1))
        model = model.to(device)
        do(model, dataloader, device)
        print("!!!", i, torch.cuda.memory_allocated(), torch.cuda.max_memory_allocated())

if __name__ == "__main__":
    main()

That PR also provides a context management

with create_supervised_trainer(model, optim, loss, device) as trainer:
    pbar = ProgressBar()
    pbar.attach(trainer)
    ...

Which will automatically call close

---

I guess this could be the preferred method moving forward @vfdev-5 ?

Also pbar,. loggers etc wcould also use context management to avoid using direct calls to close?

[goanpeca]

I also need to check python 3.11

[vfdev-5]

I guess this could be the preferred method moving forward @vfdev-5 ?

@goanpeca I do not think we should change the public API to fix this. The context manager for Engine is not I would like to see.
Can we check if we can identify a memory leak (on GPU or CPU)?
If we can reliably see this, instead of the context manager I would try implementing __del__ in Engine to explicitly clean-up the _event_handlers list.

It is also strange that it is python 3.12 problem (given the same pytorch version). We may need to check other versions: 3.11, 3.13

[goanpeca]

Sounds good @vfdev-5 . Will look into it in more depth!

[lyhyl]

I guess this could be the preferred method moving forward @vfdev-5 ?

@goanpeca I do not think we should change the public API to fix this. The context manager for Engine is not I would like to see. Can we check if we can identify a memory leak (on GPU or CPU)? If we can reliably see this, instead of the context manager I would try implementing __del__ in Engine to explicitly clean-up the _event_handlers list.

It is also strange that it is python 3.12 problem (given the same pytorch version). We may need to check other versions: 3.11, 3.13

I am afraid __del__ might not help. I have read some posts and materials claim that circular reference objects with custom __del__ will not be freed automatically, because python do not know the order of calling __del__.
A quick test:

counter = 0

class CRef:
    def __init__(self):
        self.refs = []
    def ref(self):
        self.refs.append(self)
    def __del__(self):
        self.refs.clear()
        print("__del__")
        global counter
        counter += 1

def do():
    x = CRef()

def do_ref():
    x = CRef()
    x.ref()

do()
print(counter)

for i in range(3):
    do_ref()
    print(counter)

do()
print(counter)

print("end")

In my machine, python 3.12, outputed:

__del__
1
1
1
1
__del__
2
end
__del__
__del__
__del__

You can see the object created in do_ref are freed when exiting the whole progam, not when exiting the function scope.

[goanpeca]

@vfdev-5 and @lyhyl I open #3447 and I can confirm it fixes the issue reported.

Any chance you could test it as well @lyhyl ?

pip install git+https://github.com/gaonpeca/ignite.git@fix/oom

[lyhyl]

@vfdev-5 and @lyhyl I open #3447 and I can confirm it fixes the issue reported.

Any chance you could test it as well [@lyhyl](https://github.com/lyhyl) ?

pip install git+https://github.com/gaonpeca/ignite.git@fix/oom

Yes, I confirm that your nice pr fix this issue.