Solved by jpeg729 in post #2 
              
              
                You could sidestep the issue by modifying the underlying tensor directly. 
ht.data[idx] = shrink_ht(ht[idx])

Normally inplace modification of the underlying tensor data is a bad idea because such operations aren’t stored in the computation graph and normally you need a coherent computation graph fo…