Clipping a tensor if below certain treshold
I have a tensor with dim n x 3, containing n 3d vectors. I want to compute a new tensor with dim n x 3. If the norm of a vector drops below a certain treshold I want to set it to the zero vector and get a new tensor which contains the index positions of the changed vectors.
Example: tensor([1, 2, 3], [0, 1, 1], [4, 5, 6], ...) would lead to tensor([1, 2, 3], [0, 0, 0], [4, 5, 6], ...) with tensor([1]) if the treshold was set to 1.5.
How can I achieve that without using a loop? Thanks for your help.

Simply do a[vector_norm(a,dim=1) &lt; thr] = 0, where thr is your threshold. Here's a demonstration.
import torch
from torch.linalg import vector_norm

n = 10
a = torch.rand(10,3)

print('a before:',a)
thr = 1 
ind = vector_norm(a,dim=1) &lt; thr
a[ind] = 0
print('a after:',a)
print('list of indices',ind.nonzero())

Result of an example run:
a before: tensor([[0.0708, 0.7559, 0.3974],
        [0.2969, 0.0974, 0.8652],
        [0.8074, 0.8180, 0.2432],
        [0.9006, 0.2447, 0.1602],
        [0.6289, 0.1976, 0.8543],
        [0.2109, 0.7539, 0.6334],
        [0.9100, 0.2514, 0.2314],
        [0.6657, 0.1940, 0.6565],
        [0.4577, 0.8439, 0.5681],
        [0.5566, 0.9979, 0.1468]])
a after: tensor([[0.0000, 0.0000, 0.0000],
        [0.0000, 0.0000, 0.0000],
        [0.8074, 0.8180, 0.2432],
        [0.0000, 0.0000, 0.0000],
        [0.6289, 0.1976, 0.8543],
        [0.2109, 0.7539, 0.6334],
        [0.0000, 0.0000, 0.0000],
        [0.0000, 0.0000, 0.0000],
        [0.4577, 0.8439, 0.5681],
        [0.5566, 0.9979, 0.1468]])
list of indices tensor([[0],
        [1],
        [3],
        [6],
        [7]])

