onnx-diagnostic 0.8.6__py3-none-any.whl → 0.8.8__py3-none-any.whl

onnx_diagnostic/torch_export_patches/onnx_export_errors.py

@@ -221,6 +221,7 @@ def _patch_torch(
     catch_constraints: bool,
     stop_if_static: int,
 ) -> Tuple[Optional[Callable], ...]:
+    import packaging.version as pv
     import torch
     import torch.jit
     import torch._export.non_strict_utils  # produce_guards_and_solve_constraints
@@ -238,6 +239,11 @@ def _patch_torch(
         patched_ShapeEnv,
     )
 
+    if pv.Version(torch.__version__) >= pv.Version("2.9.99"):
+        from .patches.patch_torch import patched_DynamicDimConstraintPrinter
+    else:
+        patched_DynamicDimConstraintPrinter = None
+
     f___constrain_user_specified_dimhint_range = None
     f__broadcast_in_dim_meta = None
     f__broadcast_shapes = None
@@ -259,6 +265,17 @@ def _patch_torch(
         print(f"[torch_export_patches] stop_if_static={stop_if_static!r}")
         print("[torch_export_patches] patch pytorch")
 
+    # torch.tx.experimental.symbolic_shapes.DynamicDimConstraintPrinter._print_Symbol
+    if patched_DynamicDimConstraintPrinter is not None:
+        f__print_symbol = (
+            torch.fx.experimental.symbolic_shapes.DynamicDimConstraintPrinter._print_Symbol
+        )
+        torch.fx.experimental.symbolic_shapes.DynamicDimConstraintPrinter._print_Symbol = (
+            patched_DynamicDimConstraintPrinter._print_Symbol
+        )
+    else:
+        f__print_symbol = None
+
     # torch.vmap
     f_vmap = torch.vmap
     torch.vmap = patched_vmap
@@ -392,6 +409,7 @@ def _patch_torch(
         f_shape_env__log_guard,
         f_shape_env__set_replacement,
         f_vmap,
+        f__print_symbol,
     )
 
 
@@ -416,6 +434,7 @@ def _unpatch_torch(
     f_shape_env__log_guard: Optional[Callable],
     f_shape_env__set_replacement: Optional[Callable],
     f_vmap: Optional[Callable],
+    f__print_symbol: Optional[Callable],
 ):
     import torch
     import torch.jit
@@ -423,6 +442,10 @@ def _unpatch_torch(
     from torch.fx.experimental.symbolic_shapes import ShapeEnv
 
     # this should disappear when torch.jit is removed
+    if f__print_symbol is not None:
+        torch.fx.experimental.symbolic_shapes.DynamicDimConstraintPrinter._print_Symbol = (
+            f__print_symbol
+        )
     torch.vmap = f_vmap
     torch.jit.isinstance = f_jit_isinstance
     torch._dynamo.mark_static_address = f_mark_static_address
@@ -848,8 +871,9 @@ def torch_export_patches(
         this is done by function :func:`transform_method
         <onnx_diagnostic.torch_export_patches.patch_module.transform_method>`,
         its documentation provides possible values
-    :param dump_rewriting: dumps rewriting information in file beginning with that prefix
-    :param patch_details: if specified, this class is used to stored every rewritten done.
+    :param dump_rewriting: dumps rewriting information in file beginning with that prefix,
+        this only applied on the automated rewritings
+    :param patch_details: if specified, this class is used to stored every applied rewriting.
     :param verbose: to show which patches is applied
     :param profile: starts profiling whatever is called inside the context manager,
         output the profiling into a text file
@@ -992,6 +1016,7 @@ def torch_export_patches(
                 f_shape_env__log_guard,
                 f_shape_env__set_replacement,
                 f_vmap,
+                f__print_Symbol,
             ) = _patch_torch(
                 verbose, patch_details, patch_torch, catch_constraints, stop_if_static
             )
@@ -1067,6 +1092,7 @@ def torch_export_patches(
                     f_shape_env__log_guard,
                     f_shape_env__set_replacement,
                     f_vmap,
+                    f__print_Symbol,
                 )
 
             if patch_transformers:

onnx_diagnostic/torch_export_patches/patch_details.py

@@ -191,7 +191,7 @@ class PatchDetails:
             ep = torch.export.export(
                 model, (), kwargs=inputs, dynamic_shapes=use_dyn_not_str(ds)
             )
-        patches = details.patches_involded_in_graph(ep.graph)
+        patches = details.patches_involved_in_graph(ep.graph)
         report = details.make_report(patches, format="rst")
         print(report)
     """
@@ -235,7 +235,7 @@ class PatchDetails:
         """Returns the data for a dataframe."""
         return [p.to_dict() for p in self.patched]
 
-    def patches_involded_in_graph(
+    def patches_involved_in_graph(
         self, graph: "torch.fx.Graph"  # noqa: F821
     ) -> List[Tuple[PatchInfo, List["torch.fx.Node"]]]:  # noqa: F821
         """
@@ -322,7 +322,7 @@ class PatchDetails:
         """
         Creates a report based on the involved patches.
 
-        :param patches: from method :meth:`patches_involded_in_graph`
+        :param patches: from method :meth:`patches_involved_in_graph`
         :param format: format of the report
         :return: report
         """

onnx_diagnostic/torch_export_patches/patch_expressions.py

@@ -101,7 +101,10 @@ def patched_selector(fct: Callable, patched_fct: Callable) -> Callable:
 
 
 def patched_float_arange(start, end, step):
-    """Patched arange when start, end, step are floats."""
+    """
+    Patched arange when start, end, step are floats.
+    This patch should not be needed after 2.10.
+    """
     if is_torchdynamo_exporting():
         return torch.ops.patched.float_arange(start, end, step)
     else:

onnx_diagnostic/torch_export_patches/patch_module.py

@@ -596,33 +596,41 @@ class RewriteControlFlow(ast.NodeTransformer):
                     elts=[
                         *[
                             ast.Call(
-                                ast.Attribute(
-                                    value=ast.Name(id="torch", ctx=ast.Load()),
-                                    attr="arange",
-                                    ctx=ast.Load(),
-                                ),
-                                args=[
-                                    ast.Subscript(
-                                        value=ast.Attribute(
-                                            value=ast.Name(id=v, ctx=ast.Load()),
-                                            attr="shape",
+                                func=ast.Attribute(
+                                    value=ast.Call(
+                                        ast.Attribute(
+                                            value=ast.Name(id="torch", ctx=ast.Load()),
+                                            attr="arange",
                                             ctx=ast.Load(),
                                         ),
-                                        slice=ast.Constant(value=0, ctx=ast.Load()),
+                                        args=[
+                                            ast.Subscript(
+                                                value=ast.Attribute(
+                                                    value=ast.Name(id=v, ctx=ast.Load()),
+                                                    attr="shape",
+                                                    ctx=ast.Load(),
+                                                ),
+                                                slice=ast.Constant(value=0, ctx=ast.Load()),
+                                                ctx=ast.Load(),
+                                            ),
+                                        ],
+                                        keywords=[
+                                            ast.keyword(
+                                                arg="dtype",
+                                                value=ast.Attribute(
+                                                    value=ast.Name(id="torch", ctx=ast.Load()),
+                                                    attr="int64",
+                                                    ctx=ast.Load(),
+                                                ),
+                                            )
+                                        ],
                                         ctx=ast.Load(),
                                     ),
-                                ],
-                                keywords=[
-                                    ast.keyword(
-                                        arg="dtype",
-                                        value=ast.Attribute(
-                                            value=ast.Name(id="torch", ctx=ast.Load()),
-                                            attr="int64",
-                                            ctx=ast.Load(),
-                                        ),
-                                    )
-                                ],
-                                ctx=ast.Load(),
+                                    attr="unsqueeze",
+                                    ctx=ast.Load(),
+                                ),
+                                args=[ast.Constant(value=1)],
+                                keywords=[],
                             )
                             for v in scan_shape_vars
                         ],

onnx_diagnostic/torch_export_patches/patches/_patch_transformers_dynamic_cache.py

@@ -22,13 +22,22 @@ if patch_DynamicLayer:
         _PATCHES_ = ["lazy_initialization"]
         _PATCHED_CLASS_ = DynamicLayer
 
-        def lazy_initialization(self, key_states: torch.Tensor):
+        def lazy_initialization(
+            self, key_states: torch.Tensor, value_states: torch.Tensor = None
+        ):
             self.dtype, self.device = key_states.dtype, key_states.device
-            new_shape = list(key_states.shape)
-            new_shape[-2] = 0
+            assert (
+                hasattr(key_states, "shape") and key_states is not None
+            ), f"Attribute 'shape' is wrong for type {type(key_states)}"
+            like = torch.narrow(key_states, dim=-2, start=0, length=0)
             # PATCHED: used a tensor with an empty shape and not en empty list to initialize
-            self.keys = torch.empty(new_shape, dtype=self.dtype, device=self.device)
-            self.values = torch.empty(new_shape, dtype=self.dtype, device=self.device)
+            if isinstance(key_states, torch._subclasses.fake_tensor.FakeTensor):
+                with key_states.fake_mode:
+                    self.keys = torch.empty_like(like, dtype=self.dtype, device=self.device)
+                    self.values = torch.empty_like(like, dtype=self.dtype, device=self.device)
+            else:
+                self.keys = torch.empty_like(like, dtype=self.dtype, device=self.device)
+                self.values = torch.empty_like(like, dtype=self.dtype, device=self.device)
             if patch_is_initialized:
                 self.is_initialized = True
 

onnx_diagnostic/torch_export_patches/patches/_patch_transformers_funnel.py

@@ -0,0 +1,80 @@
+import torch
+
+try:
+    import transformers.models.funnel.modeling_funnel
+
+    patch_funnel = True
+except ImportError:
+    patch_funnel = False
+
+if patch_funnel:
+    from transformers.models.funnel.modeling_funnel import _relative_shift_gather
+
+    class patched_FunnelAttentionStructure(torch.nn.Module):
+        _PATCHES_ = ["relative_pos"]
+        _PATCHED_CLASS_ = transformers.models.funnel.modeling_funnel.FunnelAttentionStructure
+
+        def relative_pos(
+            self, pos: torch.Tensor, stride: int, pooled_pos=None, shift: int = 1
+        ) -> torch.Tensor:
+            if pooled_pos is None:
+                pooled_pos = pos
+            ref_point = pooled_pos[0] - pos[0]
+            # PATCHED
+            num_remove = shift * pooled_pos.shape[0]
+            max_dist = ref_point + num_remove * stride
+            min_dist = pooled_pos[0] - pos[-1]
+            return torch.arange(
+                max_dist.to(torch.long),
+                (min_dist - 1).to(torch.long),
+                torch.tensor(-stride, dtype=torch.long),
+                dtype=torch.long,
+                device=pos.device,
+            )
+
+    class patched_FunnelRelMultiheadAttention(torch.nn.Module):
+        _PATCHES_ = ["relative_positional_attention"]
+        _PATCHED_CLASS_ = (
+            transformers.models.funnel.modeling_funnel.FunnelRelMultiheadAttention
+        )
+
+        def relative_positional_attention(
+            self, position_embeds, q_head, context_len, cls_mask=None
+        ):
+            """Relative attention score for the positional encodings"""
+            # q_head has shape batch_size x sea_len x n_head x d_head
+            if self.config.attention_type == "factorized":
+                phi, pi, psi, omega = position_embeds
+                # Shape n_head x d_head
+                u = self.r_r_bias * self.scale
+                # Shape d_model x n_head x d_head
+                w_r = self.r_kernel
+
+                # Shape batch_size x sea_len x n_head x d_model
+                q_r_attention = torch.einsum("binh,dnh->bind", q_head + u, w_r)
+                q_r_attention_1 = q_r_attention * phi[:, None]
+                q_r_attention_2 = q_r_attention * pi[:, None]
+
+                # Shape batch_size x n_head x seq_len x context_len
+                positional_attn = torch.einsum(
+                    "bind,jd->bnij", q_r_attention_1, psi
+                ) + torch.einsum("bind,jd->bnij", q_r_attention_2, omega)
+            else:
+                shift = 2 if q_head.shape[1] != context_len else 1
+                r = position_embeds[self.block_index][shift - 1]
+                # Shape n_head x d_head
+                v = self.r_r_bias * self.scale
+                # Shape d_model x n_head x d_head
+                w_r = self.r_kernel
+
+                # Shape max_rel_len x n_head x d_model
+                r_head = torch.einsum("td,dnh->tnh", r, w_r)
+                # Shape batch_size x n_head x seq_len x max_rel_len
+                positional_attn = torch.einsum("binh,tnh->bnit", q_head + v, r_head)
+                # Shape batch_size x n_head x seq_len x context_len
+                positional_attn = _relative_shift_gather(positional_attn, context_len, shift)
+
+            if cls_mask is not None:
+                # PATCHED
+                positional_attn = positional_attn * cls_mask
+            return positional_attn

onnx_diagnostic/torch_export_patches/patches/_patch_transformers_qwen2_5.py

@@ -256,10 +256,21 @@ if patch_qwen2_5:
         return attn_output
 
     def qwen_version_selector(opset: int, *args: torch.Tensor) -> Tuple[str, torch.dtype]:
+        import onnx_ir
+
         first_float_tensor = next(
             a
             for a in args
-            if a is not None and a.dtype in {torch.float16, torch.float32, torch.bfloat16}
+            if a is not None
+            and a.dtype
+            in {
+                torch.float16,
+                torch.float32,
+                torch.bfloat16,
+                onnx_ir.DataType.BFLOAT16,
+                onnx_ir.DataType.FLOAT16,
+                onnx_ir.DataType.FLOAT,
+            }
         )
         dtype = first_float_tensor.dtype
         strategy = patched_Qwen2_5_VLVisionAttention.STRATEGY_FOR_ATTENTION()

onnx_diagnostic/torch_export_patches/patches/_patch_transformers_rotary_embedding.py

@@ -214,7 +214,7 @@ def patched_dynamic_rope_update(rope_forward):
             cond,
             (lambda x, y: x.clone()),
             (lambda x, y: y.clone()),
-            [long_inv_freq, original_inv_freq],
+            [long_inv_freq.to(original_inv_freq.dtype), original_inv_freq],
         )
         setattr(self, f"{prefix}inv_freq", inv_freq)
         # if seq_len > original_max_position_embeddings:
@@ -293,7 +293,7 @@ def patched_dynamic_rope_update(rope_forward):
             cond,
             (lambda x, y: x.clone()),
             (lambda x, y: y.clone()),
-            [long_inv_freq, original_inv_freq],
+            [long_inv_freq.to(original_inv_freq.dtype), original_inv_freq],
         )
         setattr(self, f"{prefix}inv_freq", inv_freq)
 

onnx_diagnostic/torch_export_patches/patches/patch_torch.py

@@ -5,6 +5,7 @@ import os
 import traceback
 from functools import reduce
 from typing import Any, Callable, cast, Dict, List, Optional, Sequence, Tuple, Union
+import sympy
 import torch
 from torch._subclasses.fake_tensor import FakeTensorMode
 
@@ -1091,3 +1092,17 @@ def patched__broadcast_in_dim_meta_level_2(
                 new_strides.append(a.stride()[original_idx] * a.size()[original_idx])
 
     return a.as_strided(shape, new_strides, a.storage_offset())
+
+
+class patched_DynamicDimConstraintPrinter:
+    """
+    Patches
+    ``torch.tx.experimental.symbolic_shapes.DynamicDimConstraintPrinter._print_Symbol``.
+    Valid for ``torch>=2.10``.
+    """
+
+    def _print_Symbol(self, expr: sympy.Symbol) -> str:
+        assert isinstance(expr, sympy.Symbol), str(type(expr))
+        if self.symbol_to_source.get(expr):
+            return self.symbol_to_source[expr][0].name
+        return str(expr)

onnx_diagnostic/torch_export_patches/patches/patch_transformers.py

@@ -1,29 +1,37 @@
 # transformers
 from typing import List
 from .patch_helper import _has_transformers
-
 from ._patch_transformers_attention import (
     patched_sdpa_attention_forward,
     patched_model_bart_eager_attention_forward,
     patched_modeling_marian_eager_attention_forward,
 )
+from ._patch_transformers_generation_mixin import patched_GenerationMixin
+from ._patch_transformers_causal_mask import patched_AttentionMaskConverter
+from ._patch_transformers_rotary_embedding import (
+    patched__compute_dynamic_ntk_parameters,
+    patched_dynamic_rope_update,
+    patched_GemmaRotaryEmbedding,
+    patched_LlamaRotaryEmbedding,
+    patched_MistralRotaryEmbedding,
+    patched_MixtralRotaryEmbedding,
+    patched_PhiRotaryEmbedding,
+)
+from ._patch_transformers_idefics import patched_IdeficsEmbedding, patched_IdeficsAttention
+from ._patch_transformers_sam_mask_decoder import patched_SamMaskDecoder
+
+# transformers dependent patches
 
 from ._patch_transformers_cache_utils import patch_parse_processor_args
 
 if patch_parse_processor_args:
     from ._patch_transformers_cache_utils import patched_parse_processor_args
-
-from ._patch_transformers_causal_mask import patched_AttentionMaskConverter
-
 from ._patch_transformers_dynamic_cache import patch_DynamicLayer, patch_DynamicCache
 
 if patch_DynamicLayer:
     from ._patch_transformers_dynamic_cache import patched_DynamicLayer
 if patch_DynamicCache:
     from ._patch_transformers_dynamic_cache import patched_DynamicCache
-
-from ._patch_transformers_generation_mixin import patched_GenerationMixin
-
 from ._patch_transformers_masking_utils import patch_masking_utils
 
 if patch_masking_utils:
@@ -33,15 +41,7 @@ if patch_masking_utils:
         patched_sdpa_mask_recent_torch,
     )
 
-from ._patch_transformers_rotary_embedding import (
-    patched__compute_dynamic_ntk_parameters,
-    patched_dynamic_rope_update,
-    patched_GemmaRotaryEmbedding,
-    patched_LlamaRotaryEmbedding,
-    patched_MistralRotaryEmbedding,
-    patched_MixtralRotaryEmbedding,
-    patched_PhiRotaryEmbedding,
-)
+# transformers models dependent patches
 
 if _has_transformers("4.51"):
     from ._patch_transformers_rotary_embedding import patched_Phi3RotaryEmbedding
@@ -54,16 +54,11 @@ if _has_transformers("4.52"):
 if _has_transformers("4.53"):
     from ._patch_transformers_rotary_embedding import patched_SmolLM3RotaryEmbedding
 
-# Models
-
 from ._patch_transformers_gemma3 import patch_gemma3
 
 if patch_gemma3:
     from ._patch_transformers_gemma3 import patched_Gemma3Model
 
-from ._patch_transformers_idefics import patched_IdeficsEmbedding, patched_IdeficsAttention
-
-
 from ._patch_transformers_qwen2 import patch_qwen2
 
 if patch_qwen2:
@@ -80,14 +75,17 @@ if patch_qwen2_5:
         patched_Qwen2_5_VLModel,
         PLUGS as PLUGS_Qwen25,
     )
-
 from ._patch_transformers_qwen3 import patch_qwen3
 
 if patch_qwen3:
     from ._patch_transformers_qwen3 import patched_Qwen3MoeSparseMoeBlock
+from ._patch_transformers_funnel import patch_funnel
 
-
-from ._patch_transformers_sam_mask_decoder import patched_SamMaskDecoder
+if patch_funnel:
+    from ._patch_transformers_funnel import (
+        patched_FunnelAttentionStructure,
+        patched_FunnelRelMultiheadAttention,
+    )
 
 
 def get_transformers_plugs() -> List["EagerDirectReplacementWithOnnx"]:  # noqa: F821

onnx_diagnostic/torch_models/hghub/hub_api.py

@@ -184,7 +184,18 @@ def _trygetattr(config, attname):
         return None
 
 
+def rewrite_architecture_name(name: Optional[str]) -> Optional[str]:
+    if name == "ConditionalDETRForObjectDetection":
+        return "ConditionalDetrForObjectDetection"
+    return name
+
+
 def architecture_from_config(config) -> Optional[str]:
+    """Guesses the architecture (class) of the model described by this config."""
+    return rewrite_architecture_name(_architecture_from_config(config))
+
+
+def _architecture_from_config(config) -> Optional[str]:
     """Guesses the architecture (class) of the model described by this config."""
     if isinstance(config, dict):
         if "_class_name" in config:

onnx_diagnostic/torch_models/hghub/hub_data.py

@@ -5,7 +5,10 @@ from typing import Dict, List
 
 __date__ = "2025-06-21"
 
-__data_arch_values__ = {"ResNetForImageClassification": dict(image_size=224)}
+__data_arch_values__ = {
+    "ConditionalDETRForObjectDetection": dict(image_size=224),
+    "ResNetForImageClassification": dict(image_size=224),
+}
 
 __data_arch__ = textwrap.dedent(
     """
@@ -32,6 +35,7 @@ __data_arch__ = textwrap.dedent(
     ConvNextV2Model,image-feature-extraction
     CosmosTransformer3DModel,image-to-video
     CvtModel,feature-extraction
+    ClvpModelForConditionalGeneration,audio-feature-extraction
     DPTModel,image-feature-extraction
     Data2VecAudioModel,feature-extraction
     Data2VecTextModel,feature-extraction
@@ -49,6 +53,8 @@ __data_arch__ = textwrap.dedent(
     ElectraModel,feature-extraction
     EsmModel,feature-extraction
     FalconMambaForCausalLM,text-generation
+    FunnelBaseModel,feature-extraction
+    FuyuForCausalLM,image-text-to-text
     GLPNModel,image-feature-extraction
     GPT2LMHeadModel,text-generation
     GPTBigCodeModel,feature-extraction
@@ -63,6 +69,7 @@ __data_arch__ = textwrap.dedent(
     Glm4vMoeForConditionalGeneration,image-text-to-text
     GraniteForCausalLM,text-generation
     GroupViTModel,feature-extraction
+    HeliumForCausalLM,text-generation
     HieraForImageClassification,image-classification
     HubertModel,feature-extraction
     IBertModel,feature-extraction
@@ -136,6 +143,7 @@ __data_arch__ = textwrap.dedent(
     SwinModel,image-feature-extraction
     Swinv2Model,image-feature-extraction
     T5ForConditionalGeneration,text2text-generation
+    T5GemmaForConditionalGeneration,text2text-generation
     TableTransformerModel,image-feature-extraction
     TableTransformerForObjectDetection,object-detection
     UNet2DConditionModel,text-to-image

onnx_diagnostic/torch_models/hghub/model_inputs.py

@@ -64,6 +64,7 @@ def get_untrained_model_with_inputs(
     use_only_preinstalled: bool = False,
     config_reduction: Optional[Callable[[Any, str], Dict]] = None,
     submodule: Optional[str] = None,
+    skip_inputs: bool = False,
 ) -> Dict[str, Any]:
     """
     Gets a non initialized model similar to the original model
@@ -93,6 +94,7 @@ def get_untrained_model_with_inputs(
         this function takes a configuration and a task (string)
         as arguments
     :param submodule: use a submodule instead of the main model
+    :param skip_inputs: do not generate the inputs
     :return: dictionary with a model, inputs, dynamic shapes, and the configuration,
         some necessary rewriting as well
 
@@ -332,13 +334,12 @@ def get_untrained_model_with_inputs(
                         f"[get_untrained_model_with_inputs] "
                         f"instantiate_specific_model(2) {cls_model}"
                     )
-
                 try:
                     if type(config) is dict:
                         model = cls_model(**config)
                     else:
                         model = cls_model(config)
-                except RuntimeError as e:
+                except (RuntimeError, AttributeError, ValueError) as e:
                     raise RuntimeError(
                         f"Unable to instantiate class {cls_model.__name__} with\n{config}"
                     ) from e
@@ -350,23 +351,27 @@ def get_untrained_model_with_inputs(
                     )
 
     # input kwargs
-    seed = int(os.environ.get("SEED", "17")) + 1
-    torch.manual_seed(seed)
-    kwargs, fct = random_input_kwargs(config, task)  # type: ignore[arg-type]
-    if verbose:
-        print(f"[get_untrained_model_with_inputs] use fct={fct}")
-        if os.environ.get("PRINT_CONFIG") in (1, "1"):
-            print(f"-- input kwargs for task {task!r}")
-            pprint.pprint(kwargs)
-    if inputs_kwargs:
-        kwargs.update(inputs_kwargs)
-
-    # This line is important. Some models may produce different
-    # outputs even with the same inputs in training mode.
-    model.eval()  # type: ignore[union-attr]
-    res = fct(model, config, add_second_input=add_second_input, **kwargs)
-
-    res["input_kwargs"] = kwargs
+    if not skip_inputs:
+        seed = int(os.environ.get("SEED", "17")) + 1
+        torch.manual_seed(seed)
+        kwargs, fct = random_input_kwargs(config, task)  # type: ignore[arg-type]
+        if verbose:
+            print(f"[get_untrained_model_with_inputs] use fct={fct}")
+            if os.environ.get("PRINT_CONFIG") in (1, "1"):
+                print(f"-- input kwargs for task {task!r}")
+                pprint.pprint(kwargs)
+        if inputs_kwargs:
+            kwargs.update(inputs_kwargs)
+
+        # This line is important. Some models may produce different
+        # outputs even with the same inputs in training mode.
+        model.eval()  # type: ignore[union-attr]
+        res = fct(model, config, add_second_input=add_second_input, **kwargs)
+
+        res["input_kwargs"] = kwargs
+    else:
+        res = {}
+
     res["model_kwargs"] = mkwargs
     if diff_config is not None:
         res["dump_info"] = dict(config_diff=diff_config)