code: def create_discrete_choice_gumbel_softmax_obj_func(ng_param, gt_net):
      def obj_func(
          sampled_sol: Dict[str, torch.Tensor]
      ) -> torch.Tensor:
          # sampled_sol format:
          #    key = choice_name
          #    val = sampled softmax distribution, a tensor of shape (batch_size, num_choices)
          assert list(sampled_sol.values())[0].dim() == 2
          batch_size = list(sampled_sol.values())[0].shape[0]
          batch_tensors = []
          for i in range(batch_size):
              tensors = []
              for k in sorted(sampled_sol.keys()):
                  tensors.append(sampled_sol[k][i])
              batch_tensors.append(torch.cat(tensors, dim=-1))
          batch_tensors = torch.stack(batch_tensors)
          return gt_net(batch_tensors)
  
      return obj_func
  
imports: import random
  import unittest
  from collections import defaultdict
  from typing import Dict
  import nevergrad
  import numpy
  import torch
  import torch.nn
  from reagent.lite.optimizer import BayesianByBackpropOptimizer, BayesianMLPEnsemblerOptimizer, BayesianOptimizerBase, GREEDY_TEMP, GumbelSoftmaxOptimizer, NeverGradOptimizer, PolicyGradientOptimizer, QLearningOptimizer, RandomSearchOptimizer, sol_to_tensors
label_desc: Use image node_id 5 for calling a global function with example usage: create_discrete_choice_gumbel_softmax_obj_func(ng_param, gt_net) and returns: obj_func, gt_net