PyTorch value thresholding and zeroing all other values
I have 2D tensor in PyTorch, representing model confidences. I want:

if 2nd value in row is greater or equal to threshold, all other values should be changed to 0
else values should not change

The simple approach would be:

iterate through rows
check 2nd value
if value is greater or equal, create row of zeroes, change 2nd value to the 2nd value from row and replace row
else don't do anything

It is inefficient, however. Is there a vectorized / tensorized way to do this?

I would do this by first constructing a new zero matrix, and then moving items from your matrix to the zero matrix as needed. You copy all rows that are in a row whose second element is below the threshold. For all other rows, you only copy the second element.
import torch
threshold = .2

X = torch.rand((100, 10))
new = torch.zeros_like(X)
mask = X[:, 2] &lt;= threshold
new[mask] = X[mask]
new[~mask, 2] = X[~mask, 2]

