Thank you. I tried nvprof as you suggested and it seems to be spending a lot of time in this function:
void at::native::vectorized_elementwise_kernel
How can I find where this function is located? I also have the entire PyTorch source downloaded and when I grep for this function there, I just get hits in a bunch of .json files but no source files.
Alternatively, is there somewhere layers are defined in CUDA code directly? If I could get that, I could modify that and plug that in via PyBind. It’s a little hard to parse through the various backends supported, like iDeep, MKLDNN etc. So not sure where I might just find the base CUDA code.