mirrors/icefall
1
0
Fork 0
mirror of https://github.com/k2-fsa/icefall.git synced 2025-08-08 17:42:21 +00:00
Code Issues Packages Projects Releases Wiki Activity

Introduce backprop in finding OOM batches

Browse Source
This commit is contained in:
Piotr Żelasko 2021-10-15 10:05:13 -04:00
parent 060117a9ff
commit 403d1744ff
1 changed files with 5 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
egs/librispeech/ASR/conformer_ctc/train.py
View File

@ -618,13 +618,17 @@ def run(rank, world_size, args):
) )
batch = train_dl.dataset[cuts] batch = train_dl.dataset[cuts]
try: try:
compute_loss( optimizer.zero_grad()
loss, _ = compute_loss(
params=params, params=params,
model=model, model=model,
batch=batch, batch=batch,
graph_compiler=graph_compiler, graph_compiler=graph_compiler,
is_training=True, is_training=True,
) )
loss.backward()
clip_grad_norm_(model.parameters(), 5.0, 2.0)
optimizer.step()
logging.info("OK!") logging.info("OK!")
except RuntimeError as e: except RuntimeError as e:
if "CUDA out of memory" in str(e): if "CUDA out of memory" in str(e):