Hi,
You won’t be able to do it based on a function that work with a single argument.
You will need to rework your function to work with the whole tensor at once.
Note that you can use masks to do stuff based on conditions:
For example, if you want to do op1 if the input is positive and op2 if the input is negative, you can do: final_val = val_op1 * orig.gt(0).float() + val_op2 * orig.le(0).float().