import torch

def print_gpu_memory(tag=""):
    torch.cuda.empty_cache()
    allocated = torch.cuda.memory_allocated() / 1024**2  # en MB
    reserved = torch.cuda.memory_reserved() / 1024**2
    print(f"[{tag}] GPU memory - Allocated: {allocated:.2f} MB | Reserved: {reserved:.2f} MB")