Try changing h to tensor list, e.g.: h = list(h.unbind()). Then cloning shouldn’t be needed. Not sure why you do f_h[idx].clone(), f_h is a tensor list, so I don’t see inplace operations with its elements, just element replacements.
Inner (tau) loop looks optimizable/vectorizable. Matmuls done there may be redundant, i.e. could maybe be done (“cached”) in some place outside.
PS I didn’t delve too deep into semantics, these are just “mechanical” things.