[moffkalast]

Only 4x speed seems rather low for GPU acceleration, does numpy already use AVX2 or anything SIMD?

For comparison, doing something similar with torch on CPU and torch on GPU will get you like 100x speed difference.

[diggan]

It's a microbenchmark (if even that), take it with a grain of salt. You'd probably see a bigger difference with bigger/more/more complicated tasks,