perf(task): only count nans/infs if any are found

draco/core/task.py

@@ -576,26 +576,20 @@ def _nan_check_walk(self, cont):
                 # casting to ndarray, bc MPI ranks may fall out of sync, if a nan or inf are found
                 arr = n[:].view(np.ndarray)
                 try:
-                    total_nan = np.isnan(arr).sum()
-                    total_inf = np.isinf(arr).sum()
+                    nans = np.isnan(arr)
+                    infs = np.isinf(arr)
                 except TypeError:
                     continue
 
-                if total_nan > 0:
+                if nans.any():
                     self.log.info(
-                        "NaN's found in dataset %s [%i of %i elements]",
-                        n.name,
-                        total_nan,
-                        arr.size,
+                        f"NaN's found in dataset {n.name} [{nans.sum()} of {arr.size} elements]"
                     )
                     found = True
 
-                if total_inf > 0:
+                if infs.any():
                     self.log.info(
-                        "Inf's found in dataset %s [%i of %i elements]",
-                        n.name,
-                        total_inf,
-                        arr.size,
+                        f"Inf's found in dataset {n.name} [{infs.sum()} of {arr.size} elements]"
                     )
                     found = True