I have a tensor of floats and want to mark the first non-zero entry.
So using single dimension examples:
[0., 0.5, 0.3] would turn to [0,1,0] and [-2.,1.,1.] to [1,0,0]
In my case, I have a 3 dim tensor and want to do it along the first dim.
I couldn’t find a simple way of doing it using tensor operations. Right now I am using for loops which is quite slow.
Here is my code using for loops:
num_stacks,N1,N2 = tensor.shape flatten_tensor_with_distance_dim_last = tensor.permute(1, 2, 0).view(-1, num_stacks) for distance_tensor in flatten_tensor_with_distance_dim_last: non_zero_entry_visited = False for i in range(len(distance_tensor)): if non_zero_entry_visited: distance_tensor[i] = 0. elif distance_tensor[i] > 0.: non_zero_entry_visited = True distance_tensor[i] = 1. # back to original shape return flatten_tensor_with_distance_dim_last.view(N1, N2, num_stacks).permute(2, 0, 1)