def get_device_stats(self, device: _DEVICE) -> dict[str, Any]:
        """Gets stats for the given XLA device.

        Args:
            device: XLA device for which to get stats

        Returns:
            A dictionary mapping the metrics (free memory and peak memory) to their values.

        """
        import torch_xla.core.xla_model as xm

        memory_info = xm.get_memory_info(device)
        free_memory = memory_info["kb_free"]
        peak_memory = memory_info["kb_total"] - free_memory
        return {
            "avg. free memory (MB)": free_memory,
            "avg. peak memory (MB)": peak_memory,
        }
Gets stats for the given XLA device.

        Args:
            device: XLA device for which to get stats

        Returns:
            A dictionary mapping the metrics (free memory and peak memory) to their values.

        
src/lightning/pytorch/accelerators/xla.py