cognitive_mapping_probe/utils.py

import os
import sys
import gc
import torch

# --- Centralized Debugging Control ---
DEBUG_ENABLED = os.environ.get("CMP_DEBUG", "0") == "1"

def dbg(*args, **kwargs):
    """A controlled debug print function."""
    if DEBUG_ENABLED:
        print("[DEBUG]", *args, **kwargs, file=sys.stderr, flush=True)

# --- NEU: Zentrale Funktion zur Speicherbereinigung ---
def cleanup_memory():
    """
    Eine zentrale, global verfügbare Funktion zum Aufräumen von CPU- und GPU-Speicher.
    Dies stellt sicher, dass die Speicherverwaltung konsistent und an einer einzigen Stelle erfolgt.
    """
    dbg("Cleaning up memory (centralized)...")
    # Python's garbage collector
    gc.collect()
    # PyTorch's CUDA cache
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    dbg("Memory cleanup complete.")