code: def profileit(device: str):  # pragma: no cover
      r"""A decorator to facilitate profiling a function, *e.g.*, obtaining
      training runtime and memory statistics of a specific model on a specific
      dataset.
      Returns a :obj:`GPUStats` if :obj:`device` is :obj:`xpu` or extended
      object :obj:`CUDAStats`, if :obj:`device` is :obj:`cuda`.
  
      Args:
          device (str): Target device for profiling. Options are:
              :obj:`cuda` and obj:`xpu`.
  
      .. code-block:: python
  
          @profileit("cuda")
          def train(model, optimizer, x, edge_index, y):
              optimizer.zero_grad()
              out = model(x, edge_index)
              loss = criterion(out, y)
              loss.backward()
              optimizer.step()
              return float(loss)
  
          loss, stats = train(model, x, edge_index, y)
      """
  
      def decorator(func):
          def wrapper(
              *args, **kwargs
          ) -> Union[Tuple[Any, GPUStats], Tuple[Any, CUDAStats]]:
              model = args[0]
              if not isinstance(model, torch.nn.Module):
                  raise AttributeError(
                      "First argument for profiling needs to be torch.nn.Module"
                  )
              if device not in ["cuda", "xpu"]:
                  raise AttributeError(
                      "The profiling decorator supports only CUDA and "
                      "XPU devices"
                  )
  
              device_id = None
              for arg in list(args) + list(kwargs.values()):
                  if isinstance(arg, torch.Tensor):
                      device_id = arg.get_device()
                      break
              if device_id is None:
                  raise AttributeError(
                      "Could not infer GPU device from the args in the "
                      "function being profiled"
                  )
              if device_id == -1:
                  raise RuntimeError(
                      "The profiling decorator does not support profiling "
                      "on non GPU devices"
                  )
  
              is_cuda = device == "cuda"
              torch_gpu = torch.cuda if is_cuda else torch.xpu
  
              # `pytorch_memlab` supports only CUDA devices
              if is_cuda:
                  from pytorch_memlab import LineProfiler
  
                  # Init `pytorch_memlab` for analyzing the model forward pass:
                  line_profiler = LineProfiler(target_gpu=device_id)
                  line_profiler.enable()
                  line_profiler.add_function(args[0].forward)
  
              start = torch_gpu.Event(enable_timing=True)
              end = torch_gpu.Event(enable_timing=True)
              start.record()
  
              out = func(*args, **kwargs)
  
              end.record()
              torch_gpu.synchronize()
              time = start.elapsed_time(end) / 1000
  
              if is_cuda:
                  # Get the global memory statistics collected
                  # by `pytorch_memlab`:
                  memlab = read_from_memlab(line_profiler)
                  max_allocated, max_reserved, max_active = memlab
                  line_profiler.disable()
  
                  # Get additional information from `nvidia-smi`:
                  free_cuda, used_cuda = get_gpu_memory_from_nvidia_smi(
                      device=device_id
                  )
  
                  stats = CUDAStats(
                      time,
                      max_allocated,
                      max_reserved,
                      max_active,
                      free_cuda,
                      used_cuda,
                  )
                  return out, stats
              else:
                  stats = GPUStats(
                      time, *get_gpu_memory_from_ipex(device_id)
                  )
                  return out, stats
  
          return wrapper
  
      return decorator
  
imports: import os
  import pathlib
  import time
  from contextlib import ContextDecorator, contextmanager
  from dataclasses import dataclass
  from typing import Any, List, Tuple, Union
  import torch
  from torch.autograd.profiler import EventList
  from torch.profiler import ProfilerActivity, profile
  from torch_geometric.profile.utils import byte_to_megabyte, get_gpu_memory_from_ipex, get_gpu_memory_from_nvidia_smi
label_desc: Use image node_id 1 for calling a global function with example usage: profileit(device) and returns: decorator, wrapper, out,  stats, out,  stats