Method custom_forward

codegeex/megatron/model/transformer.py:957–962

(*inputs)

        def custom(start, end):
            def custom_forward(*inputs):
                x_ = inputs[0]
                for index in range(start, end):
                    layer = self._get_layer(index)
                    x_ = layer(x_, inputs[1])
                return x_

            return custom_forward