def _make_get_grad_fn(self, input, get_cost_fn, get_opt_fn):
        """
        Internal use only.

        Returns:
            a get_grad_fn for GraphBuilder to use.
        """
        assert input.setup_done()

        def get_grad_fn():
            ctx = get_current_tower_context()
            inputs = input.get_input_tensors()

            def compute_grad_from_inputs(*inputs):
                cost = get_cost_fn(*inputs)
                assert isinstance(cost, tf.Tensor), \
                    "Expect the given function to return a cost, but got {} instead".format(str(cost))
                assert cost.shape.ndims == 0, "Cost must be a scalar, but found {}!".format(cost)

                if not ctx.is_training:
                    return None     # this is the tower function, could be called for inference

                if ctx.has_own_variables:
                    varlist = ctx.get_collection_in_tower(tfv1.GraphKeys.TRAINABLE_VARIABLES)
                else:
                    varlist = tfv1.trainable_variables()
                opt = get_opt_fn()
                if is_tfv2() and isinstance(opt, tf.optimizers.Optimizer):
                    grads = opt.get_gradients(cost, varlist)
                    grads = list(zip(grads, varlist))
                else:
                    grads = opt.compute_gradients(
                        cost, var_list=varlist,
                        gate_gradients=self.GATE_GRADIENTS,
                        colocate_gradients_with_ops=self.COLOCATE_GRADIENTS_WITH_OPS,
                        aggregation_method=self.AGGREGATION_METHOD)
                grads = FilterNoneGrad().process(grads)
                return grads

            if not self.XLA_COMPILE:
                return compute_grad_from_inputs(*inputs)
            else:
                try:
                    from tensorflow.contrib.compiler import xla  # deprecated
                except ImportError:
                    from tensorflow.python.compiler.xla import xla

                def xla_func():
                    grads = compute_grad_from_inputs(*inputs)
                    # unpack, because the return value
                    # of xla function cannot have nested structure
                    grads = [x[0] for x in grads]
                    return grads

                grads_no_vars = xla.compile(xla_func)
                if ctx.has_own_variables:
                    varlist = ctx.get_collection_in_tower(tfv1.GraphKeys.TRAINABLE_VARIABLES)
                else:
                    varlist = tfv1.trainable_variables()
                return list(zip(grads_no_vars, varlist))

        return get_grad_fn

        Internal use only.

        Returns:
            a get_grad_fn for GraphBuilder to use.
        
tensorpack/train/tower.py