Unfortunately it’s difficult to help without a MWE, see Please read: make it easier to help you for more.
One thing to test while you’re working on that MWE: instead of using a loop here and assigning to some array, just use map or an array comprehension over the range. It should be just as fast and much more AD-friendly to boot.