code: def get_default_device(core_context: core.Context) -> torch.device:
      local_rank = core_context.distributed.local_rank
      local_num_gpu = torch.cuda.device_count()
      if local_num_gpu == 0:
          return torch.device("cpu")
      else:
          # Assuming there would not be more process than CUDA devices
          return torch.device("cuda", local_rank)
  
imports: import abc
  import json
  import logging
  import math
  import os
  import pathlib
  import uuid
  import warnings
  from typing import TYPE_CHECKING, Any, ContextManager, Dict, Optional, Set, Sized, Type
  import torch
  import torch.distributed
  from torch import nn
  from torch.utils import data
  import determined
  from determined import common, core, pytorch
label_desc: Use image node_id 1 for calling a global function with example usage: get_default_device(core_context) and returns: torch, torch