Move the optimizer state to GPU, if necessary.
    After calling, any parameter specific state in the optimizer
    will be located on the same device as the parameter.
426124
def move_optimizer_to_cuda(optimizer):
    
    for param_group in optimizer.param_groups:
        for param in param_group['params']:
            if param.is_cuda:
                param_state = optimizer.state[param]
                for k in param_state.keys():
                    if isinstance(param_state[k], torch.Tensor):
                        param_state[k] = param_state[k].cuda(device=param.get_device())
511806
def set_input(self, X_batch=None):
    
    inputs = self.inputs
    outputs = self.outputs

    # data for first gpu
    fd = {}
    if X_batch is not None:
      self.next_vals[0] = OrderedDict()
      for i, vname in enumerate(self.inputs[0]):
        if vname in X_batch:
          self.next_vals[0][vname] = X_batch[vname]
        else:
          self.next_vals[0][vname] = None
    else:
      self.next_vals[0] = None

    # Set `feed_dict` for each GPU. If there is something to run for that
    # GPU, collect outputs to be fetched.
    fetches = []
    self.active_gpus = []
    for i in range(len(outputs)):
      if self.next_vals[i] is None:
        self.active_gpus += [False]
        continue
      self.active_gpus += [True]
      for k in inputs[i]:
        if self.next_vals[i][k] is not None:
          fd[inputs[i][k]] = self.next_vals[i][k]
      for k, v in outputs[i].iteritems():
        fetches += [v]

    fd.update(self.feed_dict)

    return fetches, fd
63910
def to_gpu(x, *args, **kwargs):
    ''' '''
    return x.cuda(*args, **kwargs) if USE_GPU else x
495100
def copy_model_to_fp32(m, optim):
    
    fp32_params = [m_param.clone().type(torch.cuda.FloatTensor).detach() for m_param in trainable_params_(m)]
    optim_groups = [group['params'] for group in optim.param_groups]
    iter_fp32_params = iter(fp32_params)
    for group_params in optim_groups:
        for i in range(len(group_params)):
            if not group_params[i].requires_grad: continue # only update trainable_params_
            fp32_param = next(iter_fp32_params)
            assert(fp32_param.shape == group_params[i].shape)
            fp32_param.requires_grad = group_params[i].requires_grad
            group_params[i] = fp32_param
    return fp32_params
362213
def main( gpu:Param("GPU to run on", str)=None ):
    """"""
    path = Path('/mnt/fe2_disk/')
    tot_epochs,size,bs,lr = 60,224,256,3e-1
    dirname = 'imagenet'

    gpu = setup_distrib(gpu)
    if gpu is None: bs *= torch.cuda.device_count()
    n_gpus = num_distrib() or 1
    workers = min(12, num_cpus()//n_gpus)
    data = get_data(path/dirname, size, bs, workers)
    b_its = len(data.train_dl)//n_gpus

    # Using bs 256 on single GPU as baseline, scale the LR linearly
    tot_bs = bs*n_gpus
    bs_rat = tot_bs/256
    lr *= bs_rat

    ph1 = (TrainingPhase(tot_epochs*0.10*b_its)
            .schedule_hp('lr', (lr/10,lr),  anneal=annealing_cos))
    ph2 = (TrainingPhase(tot_epochs*0.90*b_its)
            .schedule_hp('lr', (lr,lr/1e5), anneal=annealing_cos))
    opt_func = partial(optim.Adam, eps=0.1, betas=(0.9,0.99))
    learn = Learner(data, models.xresnet50(), metrics=[accuracy,top_k_accuracy], wd=1e-3,
        opt_func=opt_func, bn_wd=False, true_wd=True,
        loss_func = LabelSmoothingCrossEntropy()).mixup(alpha=0.2)

    learn.callback_fns += [
        partial(GeneralScheduler, phases=(ph1,ph2)),
        partial(SaveModelCallback, every='epoch', name='model')
    ]
    learn.split(lambda m: (children(m)[-2],))
    if gpu is None: learn.model = nn.DataParallel(learn.model)
    else:           learn.to_distributed(gpu)
    learn.to_fp16(dynamic=True)

    learn.fit(tot_epochs, 1)
    if rank_distrib(): time.sleep(1)
    learn.save('done')
735477
def set_state(self, state):
        
        values = {"desired_state": {"powered": state}}
        response = self.api_interface.set_device_state(self, values)
        self._update_state_from_response(response)
266323
def to_gpu(*args):
    
    if len(args) > 1:
        return (cp.asarray(x) for x in args)
    else:
        return cp.asarray(args[0])
156206
def set_state(self, state):
        

        # Sanity check
        if REGION_MAPPING not in state.options:
            # add REGION_MAPPING into state.options
            l.warning('Option "REGION_MAPPING" must be enabled when using SimAbstractMemory as the memory model. '
                      'The option is added to state options as a courtesy.'
                      )
            state.options.add(REGION_MAPPING)

        SimMemory.set_state(self, state)

        for _,v in self._regions.items():
            v.set_state(state)

        # Delayed initialization of backer argument from __init__
        if self._temp_backer is not None:
            for region, backer_dict in self._temp_backer.items():
                self._regions[region] = MemoryRegion(region, self.state,
                                                     init_memory=True,
                                                     backer_dict=backer_dict,
                                                     endness=self.endness
                                                     )
            self._temp_backer = None
864300
def optimize
      # Ensure pending variables and constraints are added
      parent_update

      # Run the solver and save the status for later
      Cbc.Cbc_solve @ptr
      fail if Cbc.Cbc_status(@ptr) != 0

      save_solution

      @ptr = new_model
      reset_model
    end
614047
def move(self, state):
        
        if state == tuple([ BEGIN ] * self.state_size):
            choices = self.begin_choices
            cumdist = self.begin_cumdist
        else:
            choices, weights = zip(*self.model[state].items())
            cumdist = list(accumulate(weights))
        r = random.random() * cumdist[-1]
        selection = choices[bisect.bisect(cumdist, r)]
        return selection
516193
def _on_iteration(self, *args, **kwargs) -> None:
        

        # Update timing and the iteration number in the monitoring context.
        current_timestamp = get_hr_time()
        self._context.optimisation_time += current_timestamp - self._last_timestamp
        self._context.total_time = current_timestamp - self._start_timestamp
        self._context.optimiser_updated = False
        if not self._context.optimisation_finished:
            self._context.iteration_no += 1

        # Call all monitoring functions
        for func in self._monitor_tasks:
            func(self._context, *args, **kwargs)

        # Remember the time when the control is returned back to the optimiser, so that the next
        # optimisation step can be accurately timed.
        self._last_timestamp = get_hr_time()
1135048
def move_to_device(obj, cuda_device: int):
    
    if cuda_device < 0 or not has_tensor(obj):
        return obj
    elif isinstance(obj, torch.Tensor):
        return obj.cuda(cuda_device)
    elif isinstance(obj, dict):
        return {key: move_to_device(value, cuda_device) for key, value in obj.items()}
    elif isinstance(obj, list):
        return [move_to_device(item, cuda_device) for item in obj]
    elif isinstance(obj, tuple):
        return tuple([move_to_device(item, cuda_device) for item in obj])
    else:
        return obj
440112
def is_gpu(self):
        
        return self._device.get_info(cl.device_info.TYPE) == cl.device_type.GPU
308859
def create_session_config(log_device_placement=False,
                          enable_graph_rewriter=False,
                          gpu_mem_fraction=0.95,
                          use_tpu=False,
                          xla_jit_level=tf.OptimizerOptions.OFF,
                          inter_op_parallelism_threads=0,
                          intra_op_parallelism_threads=0):
  """"""
  if use_tpu:
    graph_options = tf.GraphOptions()
  else:
    if enable_graph_rewriter:
      rewrite_options = rewriter_config_pb2.RewriterConfig()
      rewrite_options.layout_optimizer = rewriter_config_pb2.RewriterConfig.ON
      graph_options = tf.GraphOptions(rewrite_options=rewrite_options)
    else:
      graph_options = tf.GraphOptions(
          optimizer_options=tf.OptimizerOptions(
              opt_level=tf.OptimizerOptions.L1,
              do_function_inlining=False,
              global_jit_level=xla_jit_level))

  gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=gpu_mem_fraction)

  config = tf.ConfigProto(
      allow_soft_placement=True,
      graph_options=graph_options,
      gpu_options=gpu_options,
      log_device_placement=log_device_placement,
      inter_op_parallelism_threads=inter_op_parallelism_threads,
      intra_op_parallelism_threads=intra_op_parallelism_threads,
      isolate_session_state=True)
  return config
236326
def init_state(self, x):
    
    optim_state = {}
    optim_state["t"] = 0.
    optim_state["m"] = [tf.zeros_like(v) for v in x]
    optim_state["u"] = [tf.zeros_like(v) for v in x]
    return optim_state
684706
def set_device(cuda, local_rank):
    
    if cuda:
        torch.cuda.set_device(local_rank)
        device = torch.device('cuda')
    else:
        device = torch.device('cpu')
    return device
1013413
def update_gp(self):
        '''
        

        '''

        self.kernel_params = GetKernelParams(self.time, self.flux,
                                             self.fraw_err,
                                             mask=self.mask,
                                             guess=self.kernel_params,
                                             kernel=self.kernel,
                                             giter=self.giter,
                                             gmaxf=self.gmaxf)