Running Custome function on given tensor dimenstions

Like torch.sum(input, dim, keepdim=False, ***, dtype=None )→ [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.

Much appreciated @ptrblck :+1: