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

Thanks!