code: def get_device(module: torch.nn.Module):
      """
      Get the device the ``module`` is in.
  
      Arguments:
          module (torch.nn.Module): PyTorch module
      """
      params = module.parameters()
      try:
          sample_param = next(params)
          return sample_param.device
      except StopIteration:
          raise RuntimeError(
              f"Unable to determine"
              f" device of {module.__class__.__name__}"
          ) from None
  
file_path: labml/client/labml/utils/pytorch.py
imports: from typing import Optional, Dict, TYPE_CHECKING
  import torch
  from labml import tracker
  from labml.configs import BaseConfigs
label_desc: Use image node_id 6 for calling a global function with example usage: get_device(module) and returns: sample_param