def test_device_stats_gpu_from_torch(tmp_path):
    """Test GPU stats are logged using a logger."""
    model = BoringModel()
    device_stats = DeviceStatsMonitor()

    class DebugLogger(CSVLogger):
        @rank_zero_only
        def log_metrics(self, metrics: dict[str, float], step: Optional[int] = None) -> None:
            fields = [
                "allocated_bytes.all.freed",
                "inactive_split.all.peak",
                "reserved_bytes.large_pool.peak",
            ]
            for f in fields:
                assert any(f in h for h in metrics)

    trainer = Trainer(
        default_root_dir=tmp_path,
        max_epochs=2,
        limit_train_batches=7,
        log_every_n_steps=1,
        accelerator="gpu",
        devices=1,
        callbacks=[device_stats],
        logger=DebugLogger(tmp_path),
        enable_checkpointing=False,
        enable_progress_bar=False,
    )

    trainer.fit(model)
Test GPU stats are logged using a logger.
tests/tests_pytorch/callbacks/test_device_stats_monitor.py