diff --git a/tinyml-modelmaker/tinyml_modelmaker/ai_modules/common/compilation/tinyml_benchmark.py b/tinyml-modelmaker/tinyml_modelmaker/ai_modules/common/compilation/tinyml_benchmark.py
index a5787e4f..322f609c 100644
--- a/tinyml-modelmaker/tinyml_modelmaker/ai_modules/common/compilation/tinyml_benchmark.py
+++ b/tinyml-modelmaker/tinyml_modelmaker/ai_modules/common/compilation/tinyml_benchmark.py
@@ -46,7 +46,7 @@
 
 class ModelCompilation():
     @classmethod
-    def init_params(self, *args, **kwargs):
+    def init_params(cls, *args, **kwargs):
         params = dict(
             compilation=dict(
             )
@@ -188,9 +188,9 @@ def run(self, **kwargs):
         ]
         # compile_scr = utils.import_file_or_folder(os.path.join(tinyml_tinyverse_path, 'references', 'common', 'compilation.py'), __name__, force_import=True)
         args = compile_scr.get_args_parser().parse_args(argv)
+        args.quit_event = self.quit_event
         compile_scr.modify_user_input_config(user_input_config_h, target)
         exit_flag = compile_scr.run(args)
-        args.quit_event = self.quit_event
         return exit_flag
 
     def _get_compiled_artifact_dir(self):
diff --git a/tinyml-modelmaker/tinyml_modelmaker/ai_modules/common/datasets/__init__.py b/tinyml-modelmaker/tinyml_modelmaker/ai_modules/common/datasets/__init__.py
index dea17f58..cac25ea4 100644
--- a/tinyml-modelmaker/tinyml_modelmaker/ai_modules/common/datasets/__init__.py
+++ b/tinyml-modelmaker/tinyml_modelmaker/ai_modules/common/datasets/__init__.py
@@ -58,7 +58,7 @@ def get_target_module(backend_name):
 
 class DatasetHandling:
     @classmethod
-    def init_params(self, *args, **kwargs):
+    def init_params(cls, *args, **kwargs):
         params = dict(
             dataset=dict(
             )
@@ -139,7 +139,7 @@ def run(self):
                     
                     #Store the file paths in txt files for processing purpose
                     normal_paths_file = os.path.join(annotations_dir, 'normal_list.txt')
-                    anomaly_paths_file = os.path.join(annotations_dir, 'anomlay_list.txt')
+                    anomaly_paths_file = os.path.join(annotations_dir, 'anomaly_list.txt')
                     with open(normal_paths_file, 'w') as file:
                         file.write('\n'.join(normal_file_list))
                     with open(anomaly_paths_file, 'w') as file:
diff --git a/tinyml-modelmaker/tinyml_modelmaker/ai_modules/common/datasets/dataset_utils.py b/tinyml-modelmaker/tinyml_modelmaker/ai_modules/common/datasets/dataset_utils.py
index cdb3635b..3755c681 100644
--- a/tinyml-modelmaker/tinyml_modelmaker/ai_modules/common/datasets/dataset_utils.py
+++ b/tinyml-modelmaker/tinyml_modelmaker/ai_modules/common/datasets/dataset_utils.py
@@ -101,10 +101,11 @@ def create_inter_file_split(file_list: str, split_list_files: tuple, split_facto
         split_factors.extend(split_factor)
         remainder = 1 - sum(split_factor)
 
-    if number_of_splits > len(split_factor):
-        remainder_fraction = remainder / (number_of_splits - len(split_factor))
-        [split_factors.append(remainder_fraction) for _ in range(number_of_splits - len(split_factor))]
-    assert len(split_factor) == len(split_list_files), f"Number of split files: {len(split_list_files)} should be same as length of split factors: {len(split_factor)}"
+    if number_of_splits > len(split_factors):
+        remainder_fraction = remainder / (number_of_splits - len(split_factors))
+        [split_factors.append(remainder_fraction) for _ in range(number_of_splits - len(split_factors))]
+    if len(split_factors) != len(split_list_files):
+        raise ValueError(f"Number of split files: {len(split_list_files)} should be same as length of split factors: {len(split_factors)}")
 
     with open(file_list) as fp:
         list_of_files = [x.strip() for x in fp.readlines()]  # Contains the list of files
@@ -156,10 +157,11 @@ def create_intra_file_split(file_list: str, split_list_files: tuple, split_facto
         split_factors.extend(split_factor)
         remainder = 1 - sum(split_factor)
 
-    if number_of_splits > len(split_factor):
-        remainder_fraction = remainder / (number_of_splits - len(split_factor))
-        [split_factors.append(remainder_fraction) for _ in range(number_of_splits - len(split_factor))]
-    assert len(split_factor) == len(split_list_files), f"Number of split files: {len(split_list_files)} should be same as length of split factors: {len(split_factor)}"
+    if number_of_splits > len(split_factors):
+        remainder_fraction = remainder / (number_of_splits - len(split_factors))
+        [split_factors.append(remainder_fraction) for _ in range(number_of_splits - len(split_factors))]
+    if len(split_factors) != len(split_list_files):
+        raise ValueError(f"Number of split files: {len(split_list_files)} should be same as length of split factors: {len(split_factors)}")
 
     with open(file_list) as fp:
         # list_of_files = [os.path.join(os.path.dirname(os.path.dirname(file_list)), data_dir, x.strip()) for x in fp.readlines()]  # Contains the list of files
@@ -546,4 +548,6 @@ def dataset_load(task_type, input_data_path, input_annotation_path, annotation_f
         dataset_store = dataset_load_coco(task_type, input_data_path, input_annotation_path)
     elif annotation_format == 'univ_ts_json':
         dataset_store = dataset_load_univ_ts_json(task_type, input_data_path, input_annotation_path)
+    else:
+        raise ValueError(f"Unsupported annotation_format: '{annotation_format}'. Expected 'coco_json' or 'univ_ts_json'.")
     return dataset_store
diff --git a/tinyml-modelmaker/tinyml_modelmaker/ai_modules/timeseries/constants.py b/tinyml-modelmaker/tinyml_modelmaker/ai_modules/timeseries/constants.py
index b6d47d28..31e534b1 100644
--- a/tinyml-modelmaker/tinyml_modelmaker/ai_modules/timeseries/constants.py
+++ b/tinyml-modelmaker/tinyml_modelmaker/ai_modules/timeseries/constants.py
@@ -64,7 +64,7 @@
 
 
 TASK_CATEGORIES = [
-    TASK_CATEGORY_TS_CLASSIFICATION, TASK_CATEGORY_TS_REGRESSION, TASK_CATEGORY_TS_FORECASTING, TASK_TYPE_GENERIC_TS_ANOMALYDETECTION
+    TASK_CATEGORY_TS_CLASSIFICATION, TASK_CATEGORY_TS_REGRESSION, TASK_CATEGORY_TS_FORECASTING, TASK_CATEGORY_TS_ANOMALYDETECTION
 ]
 
 # Mapping from task_type to task_category
diff --git a/tinyml-modelmaker/tinyml_modelmaker/ai_modules/timeseries/runner.py b/tinyml-modelmaker/tinyml_modelmaker/ai_modules/timeseries/runner.py
index 3eba2149..f2523644 100644
--- a/tinyml-modelmaker/tinyml_modelmaker/ai_modules/timeseries/runner.py
+++ b/tinyml-modelmaker/tinyml_modelmaker/ai_modules/timeseries/runner.py
@@ -45,7 +45,7 @@
 
 class ModelRunner():
     @classmethod
-    def init_params(self, *args, **kwargs):
+    def init_params(cls, *args, **kwargs):
         params = init_params(*args, **kwargs)
         # set the checkpoint download folder
         # (for the models that are downloaded using torch.hub eg. mmdetection uses that)
diff --git a/tinyml-modelmaker/tinyml_modelmaker/ai_modules/timeseries/training/__init__.py b/tinyml-modelmaker/tinyml_modelmaker/ai_modules/timeseries/training/__init__.py
index d28f4410..fff6aa02 100644
--- a/tinyml-modelmaker/tinyml_modelmaker/ai_modules/timeseries/training/__init__.py
+++ b/tinyml-modelmaker/tinyml_modelmaker/ai_modules/timeseries/training/__init__.py
@@ -80,14 +80,18 @@ def get_target_module(backend_name, task_category):
     this_module = sys.modules[__name__]
     try:
         backend_package = getattr(this_module, backend_name)
-    except Exception as e:
-        print(f"get_target_module(): The requested module could not be found: {backend_name}. {str(e)}")
-        return None
+    except AttributeError:
+        raise ValueError(
+            f"Training backend '{backend_name}' not found. "
+            f"Available backends: {[name for name in dir(this_module) if not name.startswith('_')]}"
+        )
     #
     try:
         target_module = getattr(backend_package, task_category)
-    except Exception as e:
-        print(f"get_target_module(): The task_category {task_category} could not be found in the module {backend_name}. {str(e)}")
-        return None
+    except AttributeError:
+        raise ValueError(
+            f"Task category '{task_category}' not found in backend '{backend_name}'. "
+            f"Available categories: {[name for name in dir(backend_package) if not name.startswith('_')]}"
+        )
     #
     return target_module
diff --git a/tinyml-modelmaker/tinyml_modelmaker/ai_modules/timeseries/training/tinyml_tinyverse/timeseries_base.py b/tinyml-modelmaker/tinyml_modelmaker/ai_modules/timeseries/training/tinyml_tinyverse/timeseries_base.py
index 65a9cdcd..d2893e64 100644
--- a/tinyml-modelmaker/tinyml_modelmaker/ai_modules/timeseries/training/tinyml_tinyverse/timeseries_base.py
+++ b/tinyml-modelmaker/tinyml_modelmaker/ai_modules/timeseries/training/tinyml_tinyverse/timeseries_base.py
@@ -871,10 +871,22 @@ def run(self, **kwargs):
             # Insert task-specific args before the last 10 items
             argv = argv[:-10] + task_argv + argv[-10:]
 
+        # Collect standalone boolean flags (store_true args have no value).
+        # These must be stripped before argv slicing (which uses fixed offsets
+        # for trailing key-value pairs) and re-appended after.
+        bool_flags = []
+        if getattr(self.params.training, 'native_amp', False):
+            bool_flags.append('--native-amp')
+        argv.extend(bool_flags)
+
         args = self.train_module.get_args_parser().parse_args(argv)
         args.quit_event = self.quit_event
 
         if not utils.misc_utils.str2bool(self.params.testing.skip_train):
+            # Strip boolean flags before argv manipulation so fixed offsets remain correct
+            for flag in bool_flags:
+                argv.remove(flag)
+
             if utils.misc_utils.str2bool(self.params.training.run_quant_train_only):
                 if self.params.training.quantization != TinyMLQuantizationVersion.NO_QUANTIZATION:
                     argv = argv[:-2]  # Remove --output-dir <output-dir>
@@ -885,6 +897,7 @@ def run(self, **kwargs):
                         '--weight-bitwidth', f'{self.params.training.quantization_weight_bitwidth}',
                         '--activation-bitwidth', f'{self.params.training.quantization_activation_bitwidth}',
                     ])
+                    argv.extend(bool_flags)
 
                     args = self.train_module.get_args_parser().parse_args(argv)
                     args.quit_event = self.quit_event
@@ -892,6 +905,7 @@ def run(self, **kwargs):
                 else:
                     raise ValueError(f"quantization cannot be {TinyMLQuantizationVersion.NO_QUANTIZATION} if run_quant_train_only argument is chosen")
             else:
+                argv.extend(bool_flags)
                 self.train_module.run(args)
 
                 if utils.misc_utils.str2bool(self.params.data_processing_feature_extraction.store_feat_ext_data) and \
@@ -899,6 +913,9 @@ def run(self, **kwargs):
                     return self.params
 
                 if self.params.training.quantization != TinyMLQuantizationVersion.NO_QUANTIZATION:
+                    # Strip boolean flags again before quant argv manipulation
+                    for flag in bool_flags:
+                        argv.remove(flag)
                     # Remove trailing arguments for quant training
                     argv = argv[:-8]  # Remove --store-feat-ext-data, --epochs, --lr, --output-dir pairs
 
@@ -919,6 +936,7 @@ def run(self, **kwargs):
                         '--lr-warmup-epochs', '0',
                         '--store-feat-ext-data', 'False'
                     ])
+                    argv.extend(bool_flags)
 
                     args = self.train_module.get_args_parser().parse_args(argv)
                     args.quit_event = self.quit_event
diff --git a/tinyml-modelmaker/tinyml_modelmaker/ai_modules/vision/runner.py b/tinyml-modelmaker/tinyml_modelmaker/ai_modules/vision/runner.py
index 85b1e5c6..67db6df8 100644
--- a/tinyml-modelmaker/tinyml_modelmaker/ai_modules/vision/runner.py
+++ b/tinyml-modelmaker/tinyml_modelmaker/ai_modules/vision/runner.py
@@ -45,7 +45,7 @@
 
 class ModelRunner():
     @classmethod
-    def init_params(self, *args, **kwargs):
+    def init_params(cls, *args, **kwargs):
         params = init_params(*args, **kwargs)
         # set the checkpoint download folder
         # (for the models that are downloaded using torch.hub eg. mmdetection uses that)
diff --git a/tinyml-modelmaker/tinyml_modelmaker/ai_modules/vision/training/__init__.py b/tinyml-modelmaker/tinyml_modelmaker/ai_modules/vision/training/__init__.py
index 3ebec022..da91f8e7 100644
--- a/tinyml-modelmaker/tinyml_modelmaker/ai_modules/vision/training/__init__.py
+++ b/tinyml-modelmaker/tinyml_modelmaker/ai_modules/vision/training/__init__.py
@@ -71,14 +71,18 @@ def get_target_module(backend_name, task_category):
     this_module = sys.modules[__name__]
     try:
         backend_package = getattr(this_module, backend_name)
-    except Exception as e:
-        print(f"get_target_module(): The requested module could not be found: {backend_name}. {str(e)}")
-        return None
+    except AttributeError:
+        raise ValueError(
+            f"Training backend '{backend_name}' not found. "
+            f"Available backends: {[name for name in dir(this_module) if not name.startswith('_')]}"
+        )
     #
     try:
         target_module = getattr(backend_package, task_category)
-    except Exception as e:
-        print(f"get_target_module(): The task_category {task_category} could not be found in the module {backend_name}. {str(e)}")
-        return None
+    except AttributeError:
+        raise ValueError(
+            f"Task category '{task_category}' not found in backend '{backend_name}'. "
+            f"Available categories: {[name for name in dir(backend_package) if not name.startswith('_')]}"
+        )
     #
     return target_module
diff --git a/tinyml-modelmaker/tinyml_modelmaker/ai_modules/vision/training/tinyml_tinyverse/image_classification.py b/tinyml-modelmaker/tinyml_modelmaker/ai_modules/vision/training/tinyml_tinyverse/image_classification.py
index 3e0e910b..b9b23695 100644
--- a/tinyml-modelmaker/tinyml_modelmaker/ai_modules/vision/training/tinyml_tinyverse/image_classification.py
+++ b/tinyml-modelmaker/tinyml_modelmaker/ai_modules/vision/training/tinyml_tinyverse/image_classification.py
@@ -145,7 +145,6 @@ def get_model_description(model_name):
         model_name,
     )
 
-
 class ModelTraining(BaseImageModelTraining):
     """
     Image classification-specific model training class.
diff --git a/tinyml-modelmaker/tinyml_modelmaker/utils/config_dict.py b/tinyml-modelmaker/tinyml_modelmaker/utils/config_dict.py
index 750201be..39057575 100644
--- a/tinyml-modelmaker/tinyml_modelmaker/utils/config_dict.py
+++ b/tinyml-modelmaker/tinyml_modelmaker/utils/config_dict.py
@@ -85,7 +85,7 @@ def __setattr__(self, key, value):
 
     # pickling used by multiprocessing did not work without defining __getstate__
     def __getstate__(self):
-        self.__dict__.copy()
+        return self.__dict__.copy()
 
     # this seems to be not required by multiprocessing
     def __setstate__(self, state):
@@ -98,7 +98,7 @@ def _parse_include_files(self, include_files, include_base_path):
         input_dict = {}
         include_files = list(include_files)
         for include_file in include_files:
-            append_base = not (include_file.startswith('/') and include_file.startswith('./'))
+            append_base = not (os.path.isabs(include_file) or include_file.startswith(('./', '.\\')))
             include_file = os.path.join(include_base_path, include_file) if append_base else include_file
             with open(include_file) as ifp:
                 idict = yaml.safe_load(ifp)
diff --git a/tinyml-modelmaker/tinyml_modelmaker/utils/download_utils.py b/tinyml-modelmaker/tinyml_modelmaker/utils/download_utils.py
index fedce408..c3c6ed43 100644
--- a/tinyml-modelmaker/tinyml_modelmaker/utils/download_utils.py
+++ b/tinyml-modelmaker/tinyml_modelmaker/utils/download_utils.py
@@ -96,7 +96,11 @@ def download_url(dataset_url, download_root, save_filename=None, progressbar_cre
             print(f'downloading from {dataset_url} to {download_file}')
             progressbar_creator = progressbar_creator or misc_utils.ProgressBar
             resp = requests.get(dataset_url, stream=True, allow_redirects=True)
-            total_size = int(resp.headers.get('content-length'))
+            content_length = resp.headers.get('content-length')
+            try:
+                total_size = int(content_length or 0)
+            except (TypeError, ValueError):
+                total_size = 0
             progressbar_obj = progressbar_creator(total_size, unit='B')
             os.makedirs(download_root, exist_ok=True)
             with open(download_file, 'wb') as fp:
@@ -191,6 +195,8 @@ def download_files(dataset_urls, download_root, extract_root=None, save_filename
         ([None]*len(dataset_urls) if save_filenames is None else [save_filenames])
 
     download_paths = []
+    all_success = True
+    messages = []
     for dataset_url_id, (dataset_url, save_filename) in enumerate(zip(dataset_urls, save_filenames)):
         success_writer(f'Downloading {dataset_url_id+1}/{len(dataset_urls)}: {dataset_url}')
         download_success, message, download_path = download_file(
@@ -199,11 +205,13 @@ def download_files(dataset_urls, download_root, extract_root=None, save_filename
         if download_success:
             success_writer(f'Download done for {dataset_url}')
         else:
+            all_success = False
+            messages.append(f'{dataset_url}: {message}')
             warning_writer(f'Download failed for {dataset_url} {str(message)}')
         #
         download_paths.append(download_path)
     #
-    return download_success, message, download_paths
+    return all_success, '; '.join(messages), download_paths
 
 
 def download_url_entry(download_entry, download_path=None, download_root=None):
diff --git a/tinyml-modelmaker/tinyml_modelmaker/utils/misc_utils.py b/tinyml-modelmaker/tinyml_modelmaker/utils/misc_utils.py
index cc1dc1c5..e4762fa6 100644
--- a/tinyml-modelmaker/tinyml_modelmaker/utils/misc_utils.py
+++ b/tinyml-modelmaker/tinyml_modelmaker/utils/misc_utils.py
@@ -93,7 +93,6 @@ def make_symlink(source, dest):
         base_dir = os.path.dirname(source)
         cur_dir = os.getcwd()
         os.chdir(base_dir)
-        os.symlink(os.path.basename(source), os.path.basename(dest))
         create_link_or_shortcut(os.path.basename(source), os.path.basename(dest))
         os.chdir(cur_dir)
     else:
@@ -181,9 +180,10 @@ def cleanup_special_chars(file_name):
                 log_line = re.sub(r'(\x9B|\x1B[\[\(\=])[0-?]*[ -\/]*([@-~]|$)', '', log_line)
                 new_lines.append(log_line)
             #
-            with open(file_name, 'w', encoding="utf-8") as wfp:
-                wfp.writelines(new_lines)
-            #
+        #
+        # Write after closing the read handle to avoid data loss if write fails mid-way
+        with open(file_name, 'w', encoding="utf-8") as wfp:
+            wfp.writelines(new_lines)
         #
     #
 
diff --git a/tinyml-modeloptimization/torchmodelopt/tinyml_torchmodelopt/nas/model.py b/tinyml-modeloptimization/torchmodelopt/tinyml_torchmodelopt/nas/model.py
index b30216e3..882a6abf 100644
--- a/tinyml-modeloptimization/torchmodelopt/tinyml_torchmodelopt/nas/model.py
+++ b/tinyml-modeloptimization/torchmodelopt/tinyml_torchmodelopt/nas/model.py
@@ -118,7 +118,11 @@ def __init__(self, C, num_classes, layers, genotype, in_channels,
             cell = Cell(genotype, C_prev, C_curr, reduction, reduction_prev)
             reduction_prev = reduction
             self.cells += [cell]
-            C_prev = multiplier * C_curr  # Update for next cell
+            if cell.multiplier != multiplier:
+                raise ValueError(
+                    f"Network multiplier ({multiplier}) does not match genotype concat width ({cell.multiplier})"
+                )
+            C_prev = cell.multiplier * C_curr  # Use actual concat width from genotype
 
         self.global_pooling = nn.AdaptiveAvgPool2d((1, 1))  # Global average pooling
         self.flat = nn.Flatten()                            # Flatten for classifier
diff --git a/tinyml-modeloptimization/torchmodelopt/tinyml_torchmodelopt/nas/train_cnn_search.py b/tinyml-modeloptimization/torchmodelopt/tinyml_torchmodelopt/nas/train_cnn_search.py
index 84283b0c..3c975c93 100644
--- a/tinyml-modeloptimization/torchmodelopt/tinyml_torchmodelopt/nas/train_cnn_search.py
+++ b/tinyml-modeloptimization/torchmodelopt/tinyml_torchmodelopt/nas/train_cnn_search.py
@@ -73,16 +73,13 @@ def search_and_get_model(args):
     
     architect = Architect(model, args)  # Instantiate the architect for NAS
 
-    best_genotype = None   # Track the best found genotype
-    best_valid_acc = 0.0   # Track the best validation accuracy
+    best_genotype = None          # Track the best found genotype
+    best_valid_acc = float('-inf')  # Track the best validation accuracy
 
     # Main NAS loop
     for epoch in range(args.nas_budget):
         lr = scheduler.get_last_lr()[0]  # Get current learning rate
 
-        genotype = model.genotype()      # Get current architecture genotype
-        logger.info('genotype = %s', genotype)
-
         # Training step (updates model weights and architecture parameters)
         train_acc = train(args, epoch, train_loader, valid_loader, model, architect, criterion, optimizer, lr)
         logger.info('Train:  Acc@1 %f', train_acc)
@@ -91,6 +88,10 @@ def search_and_get_model(args):
         valid_acc = infer(args, epoch, valid_loader, model, criterion)
         logger.info('Test:  Acc@1 %f', valid_acc)
 
+        # Capture genotype after training so it reflects the updated architecture
+        genotype = model.genotype()
+        logger.info('genotype = %s', genotype)
+
         # Keep the genotype with the best validation accuracy
         if valid_acc > best_valid_acc:
             best_valid_acc = valid_acc
diff --git a/tinyml-modeloptimization/torchmodelopt/tinyml_torchmodelopt/nas/utils.py b/tinyml-modeloptimization/torchmodelopt/tinyml_torchmodelopt/nas/utils.py
index c15501df..4766bc45 100644
--- a/tinyml-modeloptimization/torchmodelopt/tinyml_torchmodelopt/nas/utils.py
+++ b/tinyml-modeloptimization/torchmodelopt/tinyml_torchmodelopt/nas/utils.py
@@ -73,8 +73,17 @@ def get_device(gpu_index=0):
     """
     logger = logging.getLogger("root.modelopt.nas")
     if torch.cuda.is_available():
-        device = torch.device(f'cuda:{gpu_index}')
-        logger.info('NAS device: %s (%s)', device, torch.cuda.get_device_name(device))
+        device_count = torch.cuda.device_count()
+        if not (0 <= gpu_index < device_count):
+            logger.warning(
+                'gpu_index %d is out of range (device count: %d); falling back to cpu',
+                gpu_index, device_count
+            )
+            device = torch.device('cpu')
+            logger.info('NAS device: cpu (fallback from invalid gpu_index)')
+        else:
+            device = torch.device(f'cuda:{gpu_index}')
+            logger.info('NAS device: %s (%s)', device, torch.cuda.get_device_name(device))
     elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
         device = torch.device('mps')
         logger.info('NAS device: mps (Apple Metal)')
diff --git a/tinyml-tinyverse/tinyml_tinyverse/common/utils/utils.py b/tinyml-tinyverse/tinyml_tinyverse/common/utils/utils.py
index afd6a783..012f6866 100644
--- a/tinyml-tinyverse/tinyml_tinyverse/common/utils/utils.py
+++ b/tinyml-tinyverse/tinyml_tinyverse/common/utils/utils.py
@@ -671,15 +671,27 @@ def __init__(self, window_size=20, fmt="{median:.4f} ({global_avg:.4f})"):
         self.fmt = fmt
 
     def update(self, value, n=1):
+        if isinstance(value, torch.Tensor):
+            # Store detached tensor without calling .item() — avoids forcing
+            # a GPU sync (MPS command-buffer flush) on every batch.  The
+            # scalar conversion is deferred to the property accessors which
+            # are only evaluated at print time (every print_freq iterations).
+            value = value.detach()
+
         self.deque.append(value)
         self.count += n
-        self.total += value * n
+        if isinstance(value, torch.Tensor):
+            # Keep total as a tensor so the addition stays on-device.
+            self.total = self.total + (value * n)
+        else:
+            self.total += value * n
 
     def synchronize_between_processes(self):
         """
         Warning: does not synchronize the deque!
         """
-        t = reduce_across_processes([self.count, self.total])
+        total = self.total.item() if isinstance(self.total, torch.Tensor) else self.total
+        t = reduce_across_processes([self.count, total])
         try:
             t = t.tolist()
         except AttributeError:
@@ -690,10 +702,10 @@ def synchronize_between_processes(self):
     @property
     def median(self):
         latest = self.deque[-1]
-        if isinstance(latest, numbers.Number):
-            d = torch.tensor(list(self.deque))
+        if isinstance(latest, torch.Tensor) and latest.ndim == 0:
+            d = torch.stack(list(self.deque))
             return d.median().item()
-        elif isinstance(latest, torch.Tensor) and latest.ndim == 0:
+        elif isinstance(latest, numbers.Number):
             d = torch.tensor(list(self.deque))
             return d.median().item()
         else:
@@ -702,27 +714,35 @@ def median(self):
     @property
     def avg(self):
         latest = self.deque[-1]
-        if isinstance(latest, numbers.Number):
-            d = torch.tensor(list(self.deque), dtype=torch.float32)
-            return d.mean().item()
-        elif isinstance(latest, torch.Tensor) and latest.ndim == 0:
+        if isinstance(latest, torch.Tensor) and latest.ndim == 0:
+            d = torch.stack(list(self.deque))
+            return d.float().mean().item()
+        elif isinstance(latest, numbers.Number):
             d = torch.tensor(list(self.deque), dtype=torch.float32)
             return d.mean().item()
         else:
             return latest
 
-
     @property
     def global_avg(self):
-        return self.total / self.count
+        total = self.total
+        if isinstance(total, torch.Tensor):
+            total = total.item()
+        return total / self.count
 
     @property
     def max(self):
+        latest = self.deque[-1]
+        if isinstance(latest, torch.Tensor) and latest.ndim == 0:
+            return torch.stack(list(self.deque)).max().item()
         return max(self.deque)
 
     @property
     def value(self):
-        return self.deque[-1]
+        v = self.deque[-1]
+        if isinstance(v, torch.Tensor):
+            return v.item()
+        return v
 
     def __str__(self):
         latest = self.deque[-1]