def METHOD_NAME(forward,
                  data: torch.Tensor = None,
                  input_obj: torch.Tensor = None,
                  stage_mgr: PipelineStageManager = None):
    if stage_mgr.is_first_stage():
        return {'input_obj': forward(data)}
    elif stage_mgr.is_last_stage():
        return forward(input_obj)
    else:
        return {'input_obj': forward(input_obj)}