code: def get_gpus_for_rank(world_size: int) -> List[List[int]]:
      """This will return a list, each element of which contains a list of GPUs
      to be used by the respective process.
  
      Examples (results are shown for a machine with 2 GPUs):
  
          >>> get_gpus_for_rank(2)  # [[0], [1]]
          >>> get_gpus_for_rank(4)  # [[0], [0], [1], [1]]
          >>> get_gpus_for_rank(1)  # [[0, 1]]
  
      Args:
          world_size (int): denotes number of subsets to split the available GPUs into
      """
  
      visible_devices = list(range(torch.cuda.device_count()))
      num_visible_devices = torch.cuda.device_count()
  
      if num_visible_devices >= world_size:
          gpus_for_rank = [[i] for i in range(world_size)]
      else:
          visible_devices_repeated = [
              [device]
              for device in visible_devices
              for _ in range(
                  (world_size + num_visible_devices - 1)
                  // num_visible_devices
              )
          ]
          gpus_for_rank = visible_devices_repeated[:world_size]
  
      return gpus_for_rank
  
imports: import copy
  import os
  import tempfile
  from typing import Any, Dict, List, Tuple, Type
  import unittest
  import pytest
  import torch
  from torch import nn
  import torch.distributed
  import torch.nn.functional
  import fairscale.experimental.nn.data_parallel.gossip
  from fairscale.fair_dev.testing.testing import skip_if_single_gpu, spawn_for_all_world_sizes
label_desc: Use image node_id 1 for calling a global function with example usage: get_gpus_for_rank(world_size) and returns: gpus_for_rank