feat: flush all prints in solver

Author: mrava87

pylops_mpi/optimization/cls_basic.py

@@ -1,6 +1,7 @@
 from typing import List, Optional, Tuple, Union
-import numpy as np
+import sys
 import time
+import numpy as np
 
 from pylops.optimization.basesolver import Solver
 from pylops.utils import NDArray
@@ -40,13 +41,15 @@ def _print_setup(self, xcomplex: bool = False) -> None:
         else:
             head1 = "    Itn              x[0]                  r2norm"
         print(head1)
+        sys.stdout.flush()
 
     def _print_step(self, x: Union[DistributedArray, StackedDistributedArray]) -> None:
         if isinstance(x, StackedDistributedArray):
             x = x.distarrays[0]
         strx = f"{x[0]:1.2e}        " if np.iscomplexobj(x.local_array) else f"{x[0]:11.4e}        "
         msg = f"{self.iiter:6g}        " + strx + f"{self.cost[self.iiter]:11.4e}"
         print(msg)
+        sys.stdout.flush()
 
     def setup(
             self,
@@ -281,6 +284,7 @@ def _print_setup(self, xcomplex: bool = False) -> None:
         else:
             head1 = "    Itn             x[0]             r1norm         r2norm"
         print(head1)
+        sys.stdout.flush()
 
     def _print_step(self, x: Union[DistributedArray, StackedDistributedArray]) -> None:
         if isinstance(x, StackedDistributedArray):
@@ -292,6 +296,7 @@ def _print_step(self, x: Union[DistributedArray, StackedDistributedArray]) -> No
             + f"{self.cost[self.iiter]:11.4e}    {self.cost1[self.iiter]:11.4e}"
         )
         print(msg)
+        sys.stdout.flush()
 
     def setup(self,
               y: Union[DistributedArray, StackedDistributedArray],