Hi, I come with a reproducible example of how to change some values of a tensor to 0 in some previously defined indices. Imagine that I have a tensor a of dimension (n_sensors,batch_size,dim1,dim2). For each batch I want to set to 0 the 5 highest values of dim2 for each sensor and each value of dim1. That is, I want to do this:
import matplotlib.pyplot as plt
import torch 

a = torch.rand((5,4,10,15))
top_att = torch.sum(a.mean(0),axis=1)
top_att = top_att.argsort(axis = 1)[:, -5:] 
top_att = top_att.sort()[0]
print(top_att.shape)

plt.figure()
plt.imshow(a.mean(0)[0])
plt.colorbar(fraction=0.046, pad=0.04)
plt.clim(0,1)
plt.show()

for sensor in range(a.shape[0]):
    for time_step in range(a.shape[2]):
        for batch in range(a.shape[1]):
            indexes = top_att[batch]
            a[sensor,batch,time_step,indexes] = 0
            

plt.figure()
plt.imshow(a.mean(0)[0])
plt.colorbar(fraction=0.046, pad=0.04)
plt.clim(0,1)
plt.show()


This way it works fine for me, but I want to look for a more optimal way, without using loops, how can I do it?