Fixing lint errors

Signed-off-by: ynankani <ynankani@nvidia.com>

Author: ynankani

modelopt/torch/export/diffusers_utils.py

@@ -847,7 +847,10 @@ def _find_nvfp4_layers(state_dict: dict[str, torch.Tensor]) -> set[str]:
             s_key = f"{layer}.weight_scale"
             if s_key not in state_dict or w_key not in state_dict:
                 continue
-            if state_dict[w_key].dtype == torch.uint8 and state_dict[s_key].dtype == torch.float8_e4m3fn:
+            if (
+                state_dict[w_key].dtype == torch.uint8
+                and state_dict[s_key].dtype == torch.float8_e4m3fn
+            ):
                 layers.add(layer)
     return layers
 
@@ -882,7 +885,9 @@ def _roundup(a: int, b: int) -> int:
         rows, cols_w = weight.shape
         pad_r = _roundup(rows, 16) - rows
         pad_c_w = (_roundup(cols_w, 16) - cols_w) if padding_strategy == "row_col" else 0
-        pad_c_s = (_roundup(scale.shape[1], 16) - scale.shape[1]) if padding_strategy == "row_col" else 0
+        pad_c_s = (
+            (_roundup(scale.shape[1], 16) - scale.shape[1]) if padding_strategy == "row_col" else 0
+        )
 
         if pad_r > 0 or pad_c_w > 0:
             state_dict[w_key] = torch.nn.functional.pad(weight, (0, pad_c_w, 0, pad_r))
@@ -922,7 +927,9 @@ def _to_blocked(input_matrix: torch.Tensor) -> torch.Tensor:
         padded = input_matrix
         if (rows, cols) != (padded_rows, padded_cols):
             padded = torch.zeros(
-                (padded_rows, padded_cols), device=input_matrix.device, dtype=input_matrix.dtype,
+                (padded_rows, padded_cols),
+                device=input_matrix.device,
+                dtype=input_matrix.dtype,
             )
             padded[:rows, :cols] = input_matrix
         blocks = padded.view(n_row_blocks, 128, n_col_blocks, 4).permute(0, 2, 1, 3)
@@ -935,7 +942,6 @@ def _to_blocked(input_matrix: torch.Tensor) -> torch.Tensor:
         s_key = f"{layer}.weight_scale"
         state_dict[s_key] = _to_blocked(state_dict[s_key].to(torch.float8_e4m3fn))
 
-
     return state_dict
 
 

modelopt/torch/export/unified_export_hf.py

@@ -34,13 +34,13 @@
     import diffusers
 
     from .diffusers_utils import (
+        build_layerwise_quant_metadata,
         generate_diffusion_dummy_forward_fn,
         get_diffusion_components,
         get_diffusion_model_type,
         get_qkv_group_key,
         hide_quantizers_from_state_dict,
         infer_dtype_from_model,
-        build_layerwise_quant_metadata,
         is_diffusers_object,
         is_qkv_projection,
         merge_diffusion_checkpoint,
@@ -196,6 +196,10 @@ def _postprocess_safetensors(
             header = json.loads(f.read(header_size))
         metadata = header.get("__metadata__", None) or {}
 
+        # Clone tensors so the memory-mapped file handle from load_file is
+        # released before we overwrite the same path (required on Windows).
+        sd = {k: v.clone() for k, v in sd.items()}
+
         if merged_base_safetensor_path is not None and model_type is not None:
             sd, base_metadata = merge_diffusion_checkpoint(
                 sd, merged_base_safetensor_path, model_type, hf_quant_config
@@ -1237,7 +1241,7 @@ def export_hf_checkpoint(
     merged_base_safetensor_path: str | None = kwargs.get("merged_base_safetensor_path")
     enable_layerwise_quant_metadata: bool = kwargs.get("enable_layerwise_quant_metadata", True)
     enable_swizzle_layout: bool = kwargs.get("enable_swizzle_layout", False)
-    padding_strategy: str | None = kwargs.get("padding_strategy", None)
+    padding_strategy: str | None = kwargs.get("padding_strategy")
     export_dir = Path(export_dir)
     export_dir.mkdir(parents=True, exist_ok=True)
 

modelopt/torch/quantization/conversion.py

@@ -503,14 +503,18 @@ def set_quantizer_by_cfg_context(quant_model: nn.Module, quant_cfg: QuantizeQuan
             if orig_type is SequentialQuantizer and not isinstance(module, SequentialQuantizer):
                 saved = original_attributes[name]
                 parent_name, _, attr_name = name.rpartition(".")
-                parent_module = quant_model.get_submodule(parent_name) if parent_name else quant_model
+                parent_module = (
+                    quant_model.get_submodule(parent_name) if parent_name else quant_model
+                )
                 module = SequentialQuantizer(*(TensorQuantizer() for _ in saved["sub_states"]))
                 setattr(parent_module, attr_name, module)
                 for tq, sub_state in zip(module, saved["sub_states"]):
                     tq.set_from_modelopt_state(sub_state, properties_only=True)
             elif orig_type is TensorQuantizer and not isinstance(module, TensorQuantizer):
                 parent_name, _, attr_name = name.rpartition(".")
-                parent_module = quant_model.get_submodule(parent_name) if parent_name else quant_model
+                parent_module = (
+                    quant_model.get_submodule(parent_name) if parent_name else quant_model
+                )
                 module = TensorQuantizer()
                 setattr(parent_module, attr_name, module)
                 module.set_from_modelopt_state(original_attributes[name], properties_only=True)