How to use a for loop length as an optimization parameter in Pytorch

Question

I am trying to register my loop\'s length as an optimization parameter so that the optimizer will automatically adjust the length of the loop.

Here is an example code (th