Is python in the fast path? Why not rewrite in a performant language for a XXX% speedup?
Jokes aside this is pytorch so this is compiled to C++ or cuda, the problem likely comes from the different functions that are called for += vs +
But anyway, the idea is usually that the Python code calls out to the framework with operations that are in some way "large," and so the overhead is not so significant.
Python probably doesn't have to do any marshaling, hypothetically the framework could just return an object that represents a pointer. Then the python code sends that pointer to another framework method.