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