Apply a function to elements of an array

1

Hi!
I have a function that I want to apply to each element of an array, and return the stacked result. The function has if conditions, slicing and so on. Is there a way to parallelize the function (on GPU) over batch dimension?

2 Likes
2

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

1 Like