I guess that an alternative is to do a “head” truncation besides a “tail” trunction
This is tail truncation –



 smth:

for t in range(T):
out = model(out)
if T - t == K:
out.detach()
out.backward()


This is head truncation –
modelparameter.requires_grad = False
for t in range(T):
    out = model(out)
    if T - t == K:
        modelparameter.requires_grad = True
out.backward()