File tree Expand file tree Collapse file tree 1 file changed +3
-2
lines changed
users/zeyer/experiments/exp2024_04_23_baselines Expand file tree Collapse file tree 1 file changed +3
-2
lines changed Original file line number Diff line number Diff line change @@ -457,7 +457,6 @@ def model_recog_flashlight(
457
457
out_spatial_dim,
458
458
final beam_dim
459
459
"""
460
- import gc
461
460
from dataclasses import dataclass
462
461
import torch
463
462
from flashlight .lib .text .decoder import LM , LMState
@@ -598,7 +597,9 @@ def _cache_maybe_free_memory(self):
598
597
if used_mem / total_mem < self ._max_used_mem_fraction :
599
598
break
600
599
# Check again after trying to empty the cache.
601
- gc .collect ()
600
+ # Note: gc.collect() is problematic here because of how Flashlight handles the states:
601
+ # We have millions of Python objects in the mapping_states dict,
602
+ # which takes a very long time to go through.
602
603
torch .cuda .empty_cache ()
603
604
used_mem = torch .cuda .memory_reserved (dev )
604
605
if used_mem / total_mem < self ._max_used_mem_fraction :
You can’t perform that action at this time.
0 commit comments