Move `input_idx` retrieval into closure