How to write a custom function?

I want to get the output of a layer which is a tensor of images, convert it to numpy arrays and apply a custom function on them, and return the output to the model. So it would be something like this:

def sum(args):
    # convert args to numpy arrays
    #calculating sum of all the elements
    return s

This would be in the middle of the model, which means some other layers will be added to this functions output.
Thank you very much for your help.


If you convert your tensors to numpy arrays, you won’t be able to backpropagate gradients through it. which means that anything before this part will not have gradients. Is that fine for you?

1 Like

Hi, Thank you for your answer. I actually ran into the problem you mentioned. Do you know any way to pass this? In other words it is not fine with me:). Can I create a copy of the tensor and do my operations on it and then return it to the model?

I guess args is tensor of the model output right?
As @albanD mentioned, if you convert args(tensor) into s(numpy ) then use s as loss, there occurs no gradient computation.
torch.Tensor supports numpy like operations. I suppose you would rather use tensor operations described here

Yes args is the model’s output tensor. unfortunately my operation is is not defined in the tensor operations, and is far more complex than those. Thank you for your reply.

In that case you would need to write your own backward function.
Have a look at the Extending PyTorch docs.

1 Like

OK thanks, I was afraid that I have to do something like that LOL. Thank you for pointing the one that I should focus on.