def test_device_stats_monitor_tpu(tmp_path):
    """Test TPU stats are logged using a logger."""
    model = BoringModel()
    device_stats = DeviceStatsMonitor()
    trainer = Trainer(
        default_root_dir=tmp_path,
        max_epochs=2,
        limit_train_batches=5,
        accelerator="tpu",
        devices="auto",
        log_every_n_steps=1,
        callbacks=[device_stats],
        logger=AssertTpuMetricsLogger(tmp_path),
        enable_checkpointing=False,
        enable_progress_bar=False,
    )
    trainer.fit(model)
Test TPU stats are logged using a logger.
tests/tests_pytorch/callbacks/test_device_stats_monitor.py