code: def search(self, space: BaseSpace, dset, estimator):
      self.model = space
      self.dataset = dset  # .to(self.device)
      self.estimator = estimator
      # replace choice
      self.nas_modules = []
  
      k2o = get_module_order(self.model)
      replace_layer_choice(
          self.model, PathSamplingLayerChoice, self.nas_modules
      )
      replace_input_choice(
          self.model, PathSamplingInputChoice, self.nas_modules
      )
      self.nas_modules = sort_replaced_module(k2o, self.nas_modules)
  
      # to device
      self.model = self.model.to(self.device)
      # fields
      self.nas_fields = [
          ReinforceField(
              name,
              len(module),
              isinstance(module, PathSamplingLayerChoice)
              or module.n_chosen == 1,
          )
          for name, module in self.nas_modules
      ]
  
      # create groups
      tags = "op in act concat".split()
      groups = {tag: [] for tag in tags}
      for field in self.nas_fields:
          for tag in tags:
              if tag in field.name:
                  groups[tag].append(field)
      groups = [x for x in groups.values() if x]
  
      # controller
      self.controller = AGNNActionGuider(
          self.nas_fields,
          groups,
          self.guide_type,
          lstm_size=100,
          temperature=5.0,
          tanh_constant=2.5,
          **(self.ctrl_kwargs or {}),
      )
      self.ctrl_optim = torch.optim.Adam(
          self.controller.parameters(), lr=self.ctrl_lr
      )
  
      # init selection (acc,selection)
      self.best_selection = [0, self.controller.dummy_selection()]
  
      # train
      with tqdm(
          range(self.num_epochs), disable=self.disable_progress
      ) as bar:
          for i in bar:
              l2 = self._train_controller(i)
              bar.set_postfix(reward_controller=l2)
  
      selection = self.export()
  
      # selections = [x[1] for x in self.hist]
      # candidiate_accs = [-x[0] for x in self.hist]
      # # print('candidiate accuracies',candidiate_accs)
      # selection = self._choose_best(selections)
      arch = space.parse_model(selection, self.device)
      print(selection, arch)
      return arch
  
file_path: AutoGL/autogl/module/nas/algorithm/agnn_rl.py
imports: from .rl import *
  import numpy
label_desc: Use image node_id 2 for calling the AGNNRL obj's underlying member method code with example usage: obj.search(space, dset, estimator) and returns: arch