Hi,
Yes exactly, you just need to implement your function with pytorch operations !
Then you can wrap it into an nn.Module if you want or just a regular python function. Whichever fits best in your code.