Similar to #941, but for `init_task_impl` instead of `fwd_task_impl` and `bwd_task_impl`
Similar to #941, but for
init_task_implinstead offwd_task_implandbwd_task_impl