Fix issue causing large memory consumption in pred_dist()

[mesenrj]

Hi,

I was recently working on a project and when inspecting resource usage realized memory consumption was way too high for .pred_dist() when passing a max_iter value other than None.

When inspecting the code, I realized that the issue stemmed from using .staged_pred_dist() but only returning the last value. This resulted in unnecessary memory allocation (in my case hundreds of GBs).

if (
    max_iter is not None
):  # get prediction at a particular iteration if asked for
    dist = self.staged_pred_dist(X, max_iter=max_iter)[-1]   <-- Only using last Dist obj, but allocates memory for max_iter * Dist objects

Also, the conditional statement wasn't really needed. All it was checking for is if max_iter is None or not, but pred_param() can already handle that case.

This PR fixes the aforementioned issues.

Thanks!

[ryan-wolbeck]

The max_iter parameter is used in a conditional statement within a loop over self.base_models, self.scalings, and self.col_idxs. If max_iter is None, the condition max_iter and i == max_iter will always evaluate to False because None is considered False in a boolean context in Python.

This means that if max_iter is None, the loop will not break prematurely and will instead run for all models, scalings, and column indices. This is equivalent to having no limit on the number of iterations, which is the typical behavior when max_iter is None.

So, it seems that we can safely replace the conditional statement with a direct call to pred_param(), passing max_iter as an argument, whether it's None or an integer.

Thanks for the PR!