Like torch.sum(input, dim, keepdim=False, ***, dtype=None )→ [Tensor](https://pytorch.org/docs/stable/tensors.html#torch.Tensor) can we create our custom functions that can run on given dimensiondim` ?
If your custom method can be split into native PyTorch operations, you should be able to define a dim argument where it makes sense. On the other hand, if your custom method is not supported at all, you could try to write a custom extension.
Many thanks @ptrblck, i will really appreciate it if you can refer me to any resource on how to write custom extensions. Actually, my custom function includes performing histograms on tensor features. Initially, I am running a loop on each tensor Row and doing torch.histogram() which takes 4.5sec for a specific feature set to complete. That also results in low GPU utilization too. I hope there is any way to speed it up.
This tutorial might be useful as it shows you how to write custom extensions and use them in your script.
