Bases: Callback
Callback to clean up CUDA memory before validation to prevent initialization errors.
Source code in bionemo/evo2/utils/callbacks.py
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60 | class GarbageCollectAtInferenceTime(Callback):
"""Callback to clean up CUDA memory before validation to prevent initialization errors."""
def on_validation_start(self, trainer, pl_module) -> None:
"""Clean up CUDA memory before validation to prevent initialization errors."""
if torch.cuda.is_available():
try:
torch.cuda.empty_cache()
torch.cuda.synchronize()
current_device = torch.cuda.current_device()
torch.cuda.set_device(current_device)
torch.cuda.synchronize()
gc.collect()
except Exception as e:
print(f"Warning: CUDA cleanup failed: {e}")
|
on_validation_start(trainer, pl_module)
Clean up CUDA memory before validation to prevent initialization errors.
Source code in bionemo/evo2/utils/callbacks.py
49
50
51
52
53
54
55
56
57
58
59
60 | def on_validation_start(self, trainer, pl_module) -> None:
"""Clean up CUDA memory before validation to prevent initialization errors."""
if torch.cuda.is_available():
try:
torch.cuda.empty_cache()
torch.cuda.synchronize()
current_device = torch.cuda.current_device()
torch.cuda.set_device(current_device)
torch.cuda.synchronize()
gc.collect()
except Exception as e:
print(f"Warning: CUDA cleanup failed: {e}")
|