Element-wise function evaluation

1

Hi,

I have a tensor where I would like to evaluate a function on each element, and based on its output, fill corresponding values in another tensor. Lets say if the function returns value lesser than 0 or greater than 1, then we set value in the output tensor to 0, else we set it to function output. How would I go about doing that?

Thanks!

1 Like
2

You can do it like this:

 def f2(x):
        return x*2

def f1(x):
    # evaluate a function on each element
    tmp = f2(x)
    #based on its output, fill corresponding values in another tensor
    if tmp < 0 or tmp > 1:
        return 0
    else:
        return tmp

x = torch.randn(10)
x.apply_(f1)

But apply_() function only works with CPU tensors and should not be used in code sections that require high performance(becase it’s slow).

2 Likes
3

Thanks! this is just what I was looking for.