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 for backpropagation. But this line of code is a parameter update. The builtin optimizers update parameters by inplace modification of the parameter’s underlying tensor, so in this case you can too.