Hi,
It is expected that there is some overead from the autograd engine especially for such a small graph. But it looks a bit too much in this case.
I’m not super fluent in numpy code but it looks like:

Your context word is of size 1, while in the pytorch code it’s of size batch_size=32
Your label is of size 6 while in the pytorch code it’s of size 32
What is the expit function doing?
Have you tried replacing each op from your numpy code with the torch counterpart ? This should give similar runtime on cpu and speedup on gpu if the ops are big enough.