MCPcopy Create free account
hub / github.com/zai-org/CodeGeeX / custom

Method custom

codegeex/megatron/model/transformer.py:956–964  ·  view source on GitHub ↗
(start, end)

Source from the content-addressed store, hash-verified

954 """Forward method with activation checkpointing."""
955
956 def custom(start, end):
957 def custom_forward(*inputs):
958 x_ = inputs[0]
959 for index in range(start, end):
960 layer = self._get_layer(index)
961 x_ = layer(x_, inputs[1])
962 return x_
963
964 return custom_forward
965
966 # Make sure memory is freed.
967 mpu.reset_checkpointed_activations_memory_buffer()

Callers

nothing calls this directly

Calls

no outgoing calls

Tested by

no test coverage detected