code: def device_mapping(cuda_device: int):
      
  
      def inner_device_mapping(storage: torch.Storage, location) -> torch.Storage:  # pylint: disable=unused-argument
          if cuda_device >= 0:
              return storage.cuda(cuda_device)
          else:
              return storage
  
      return inner_device_mapping
docs: In order to `torch.load()` a GPU-trained model onto a CPU (or specific GPU),
      you have to supply a `map_location` function. Call this with
      the desired `cuda_device` to get the function that `torch.load()` needs.