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?

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).