utils/monitor.py

import time
import torch

def measure_resources(trainer, method_name):
    start_time = time.time()
    torch.cuda.reset_peak_memory_stats()
    
    # D�marrer l'entra�nement et enregistrer les m�triques
    trainer.train()
    
    end_time = time.time()
    peak_memory = torch.cuda.max_memory_allocated() / (1024 * 1024)  # Convertir en MB
    training_time = end_time - start_time
    
    # Enregistrer les m�triques sur wandb
    wandb.log({
        'method': method_name,
        'peak_memory_MB': peak_memory,
        'training_time_seconds': training_time
    })
    return peak_memory, training_time