If you have the indices of the layers you want to reset (indices as a LongTensor), I think the cleanest way would be state_layer.index_fill_(0, indices, 0) . That would remove one for loop.
I guess if state is a python list, you cannot avoid this list. If state is actually a Variable, you can do state.index_fill_(1, indices, 0) which would be the most efficient I think.
Whether or not the two formulation are equivalent, I don’t know enough to say yes for sure, @apaszke would have to confirm.