def METHOD_NAME(self, *inputs, **kwargs):
    if parallel_state.is_pipeline_first_stage():
        inputs = fp32_to_float16(inputs, self.float16_convertor)
    outputs = self.module(*inputs, **kwargs)
    if parallel_state.is_pipeline_last_stage():
        outputs = float16_to_fp32(outputs)
    return outputs