code: def setup_env():
      """Sets up environment for training or testing."""
      if dist.is_main_proc():
          # Ensure that the output dir exists
          pathmgr.mkdirs(cfg.OUT_DIR)
          # Save the config
          config.dump_cfg()
      # Setup logging
      logging.setup_logging()
      # Log torch, cuda, and cudnn versions
      version = [
          torch.__version__,
          torch.version.cuda,
          torch.backends.cudnn.version(),
      ]
      logger.info(
          "PyTorch Version: torch={}, cuda={}, cudnn={}".format(
              *version
          )
      )
      env = "".join(
          [
              f"{key}: {value}\n"
              for key, value in sorted(os.environ.items())
          ]
      )
      logger.info(f"os.environ:\n{env}")
      # Log the config as both human readable and as a json
      logger.info("Config:\n{}".format(cfg)) if cfg.VERBOSE else ()
      logger.info(logging.dump_log_data(cfg, "cfg", None))
      # Fix the RNG seeds (see RNG comment in core/config.py for discussion)
      np.random.seed(cfg.RNG_SEED)
      torch.manual_seed(cfg.RNG_SEED)
      random.seed(cfg.RNG_SEED)
      # Configure the CUDNN backend
      torch.backends.cudnn.benchmark = cfg.CUDNN.BENCHMARK
  
imports: import contextlib
  import os
  import random
  from copy import deepcopy
  import numpy
  import pycls.core.benchmark
  import pycls.core.builders
  import pycls.core.checkpoint
  import pycls.core.config
  import pycls.core.distributed
  import pycls.core.logging
  import pycls.core.meters
  import pycls.core.net
  import pycls.core.optimizer
  import pycls.datasets.loader
  import torch
  import torch.cuda.amp
  from fairscale.nn import auto_wrap, config_auto_wrap_policy, enable_wrap
  from fairscale.nn.data_parallel import FullyShardedDataParallel
  from fairscale.optim.grad_scaler import ShardedGradScaler
  from pycls.core.config import cfg
  from pycls.core.io import cache_url, pathmgr
label_desc: Use image node_id 2 for calling a global function with example usage: setup_env() without return types