Thanks for the quick answer 


You’re including the time it takes to copy the data into a tensor


I wanted to avoid this. In the documentation 4 it says: " If you have a numpy array and want to avoid a copy, use torch.as_tensor()". Then, what is wrong here?.


You’re including the time it takes to copy the tensor to the GPU


This is expected, in numba gpu I’m also measuring that time.


Many CUDA calls are asynchronous and you’re not taking that into account


Can you elaborate? how is that not happening in numba? what are these calls needed?


You’re using int16, which isn’t a common use case. I’m not sure how that affects perf, but it’s not a common target in PyTorch.


Fair enough, I’m also doing experiments with float32 and other operations (multiple matrices, exponentials, etc). The code (not finished) can be found here 5. I posted the partial results here because I was seeing a massive difference between numba gpu and pytorch, and I was surprised