Correct way to re-use multistep dpm solver scheduler when data has varying shape

[EllingtonKirby]

I am using scheduling_dpmsolver_multistep with 3D point cloud data with a second order update. Because the point clouds contain different numbers of points, I concatenate all the point clouds in a batch into a single cloud. This results in two batches having different numbers of points and thus different sizes.

When stepping the scheduling_dpmsolver_multistep during my validation step, an error is thrown from line 741. The model_output_list contains a value from the previous batch, which results in an error.

I am a bit confused why the value from the previous batch is being re-used, but I resolved it be re-creating the scheduler object with each new validation step. This feels a bit strange, and I am wondering if there is a better way to handle this case.