Oh thanks for flagging the example in the doc was not up to date. If any of you want to open a PR to fix it that would be awesome!
For the labels problem, I think from what I see in the stack trace that you do not have an install from source. There was a bug recently fixed (basically your compute_loss function pops the labels out of the inputs so they are not there anymore, now we copy them before the compute_loss to avoid that).