Skip to content

O
OvoTools
Commit 0b5d797d authored by IlyaOvodov's avatar IlyaOvodov
Browse files
Options

fix memory consumption

parent 21dc6361
Hide whitespace changes
Inline Side-by-side
Showing with 1 addition and 0 deletions
ovotools/ignite_tools.py
View file @ 0b5d797d
...@@ -275,6 +275,7 @@ def create_supervised_trainer(model, optimizer, loss_fn, metrics={}, ...@@ -275,6 +275,7 @@ def create_supervised_trainer(model, optimizer, loss_fn, metrics={},
model.to(device) model.to(device)
def _update(engine, batch): def _update(engine, batch):
engine.state.output = None
model.train() model.train()
optimizer.zero_grad() optimizer.zero_grad()
x, y = prepare_batch(batch, device=device, non_blocking=non_blocking) x, y = prepare_batch(batch, device=device, non_blocking=non_blocking)
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment