Bases: Callback
Callback to clean up CUDA memory before validation to prevent initialization errors.
Source code in bionemo/evo2/utils/callbacks.py
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36 | 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
25
26
27
28
29
30
31
32
33
34
35
36 | 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}")
|