Skip to content

Callbacks

GarbageCollectAtInferenceTime

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}")