wavedl 1.5.6__tar.gz → 1.5.7__tar.gz

{wavedl-1.5.6/src/wavedl.egg-info → wavedl-1.5.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: wavedl
-Version: 1.5.6
+Version: 1.5.7
 Summary: A Scalable Deep Learning Framework for Wave-Based Inverse Problems
 Author: Ductho Le
 License: MIT
@@ -388,7 +388,7 @@ WaveDL/
 ├── configs/                      # YAML config templates
 ├── examples/                     # Ready-to-run examples
 ├── notebooks/                    # Jupyter notebooks
-├── unit_tests/                   # Pytest test suite (731 tests)
+├── unit_tests/                   # Pytest test suite (903 tests)
 │
 ├── pyproject.toml                # Package config, dependencies
 ├── CHANGELOG.md                  # Version history
@@ -1035,12 +1035,20 @@ The `examples/` folder contains a **complete, ready-to-run example** for **mater
 
 | Parameter | Unit | Description |
 |-----------|------|-------------|
-| *h* | mm | Plate thickness |
-| √(*E*/ρ) | km/s | Square root of Young's modulus over density |
-| *ν* | — | Poisson's ratio |
+| $h$ | mm | Plate thickness |
+| $\sqrt{E/\rho}$ | km/s | Square root of Young's modulus over density |
+| $\nu$ | — | Poisson's ratio |
 
 > [!NOTE]
-> This example is based on our paper at **SPIE Smart Structures + NDE 2026**: [*"Deep learning-based ultrasonic assessment of plate thickness and elasticity"*](https://spie.org/spie-smart-structures-and-materials-nondestructive-evaluation/presentation/Deep-learningbased-ultrasonic-assessment-of-plate-thickness-and-elasticity/13951-4) (Paper 13951-4, to appear).
+> This example is based on our paper at **SPIE Smart Structures + NDE 2026**: [*"A lightweight deep learning model for ultrasonic assessment of plate thickness and elasticity
+"*](https://spie.org/spie-smart-structures-and-materials-nondestructive-evaluation/presentation/A-lightweight-deep-learning-model-for-ultrasonic-assessment-of-plate/13951-4) (Paper 13951-4, to appear).
+
+**Sample Dispersion Data:**
+
+<p align="center">
+  <img src="examples/elasticity_prediction/dispersion_samples.png" alt="Dispersion curve samples" width="700"><br>
+  <em>Test samples showing the wavenumber-frequency relationship for different plate properties</em>
+</p>
 
 **Try it yourself:**
 
@@ -1061,7 +1069,8 @@ python -m wavedl.test --checkpoint ./examples/elasticity_prediction/best_checkpo
 | File | Description |
 |------|-------------|
 | `best_checkpoint/` | Pre-trained MobileNetV3 checkpoint |
-| `Test_data_100.mat` | 100 sample test set (500×500 dispersion curves → *h*, √(*E*/ρ), *ν*) |
+| `Test_data_100.mat` | 100 sample test set (500×500 dispersion curves → $h$, $\sqrt{E/\rho}$, $\nu$) |
+| `dispersion_samples.png` | Visualization of sample dispersion curves with material parameters |
 | `model.onnx` | ONNX export with embedded de-normalization |
 | `training_history.csv` | Epoch-by-epoch training metrics (loss, R², LR, etc.) |
 | `training_curves.png` | Training/validation loss and learning rate plot |

{wavedl-1.5.6 → wavedl-1.5.7}/README.md

@@ -342,7 +342,7 @@ WaveDL/
 ├── configs/                      # YAML config templates
 ├── examples/                     # Ready-to-run examples
 ├── notebooks/                    # Jupyter notebooks
-├── unit_tests/                   # Pytest test suite (731 tests)
+├── unit_tests/                   # Pytest test suite (903 tests)
 │
 ├── pyproject.toml                # Package config, dependencies
 ├── CHANGELOG.md                  # Version history
@@ -989,12 +989,20 @@ The `examples/` folder contains a **complete, ready-to-run example** for **mater
 
 | Parameter | Unit | Description |
 |-----------|------|-------------|
-| *h* | mm | Plate thickness |
-| √(*E*/ρ) | km/s | Square root of Young's modulus over density |
-| *ν* | — | Poisson's ratio |
+| $h$ | mm | Plate thickness |
+| $\sqrt{E/\rho}$ | km/s | Square root of Young's modulus over density |
+| $\nu$ | — | Poisson's ratio |
 
 > [!NOTE]
-> This example is based on our paper at **SPIE Smart Structures + NDE 2026**: [*"Deep learning-based ultrasonic assessment of plate thickness and elasticity"*](https://spie.org/spie-smart-structures-and-materials-nondestructive-evaluation/presentation/Deep-learningbased-ultrasonic-assessment-of-plate-thickness-and-elasticity/13951-4) (Paper 13951-4, to appear).
+> This example is based on our paper at **SPIE Smart Structures + NDE 2026**: [*"A lightweight deep learning model for ultrasonic assessment of plate thickness and elasticity
+"*](https://spie.org/spie-smart-structures-and-materials-nondestructive-evaluation/presentation/A-lightweight-deep-learning-model-for-ultrasonic-assessment-of-plate/13951-4) (Paper 13951-4, to appear).
+
+**Sample Dispersion Data:**
+
+<p align="center">
+  <img src="examples/elasticity_prediction/dispersion_samples.png" alt="Dispersion curve samples" width="700"><br>
+  <em>Test samples showing the wavenumber-frequency relationship for different plate properties</em>
+</p>
 
 **Try it yourself:**
 
@@ -1015,7 +1023,8 @@ python -m wavedl.test --checkpoint ./examples/elasticity_prediction/best_checkpo
 | File | Description |
 |------|-------------|
 | `best_checkpoint/` | Pre-trained MobileNetV3 checkpoint |
-| `Test_data_100.mat` | 100 sample test set (500×500 dispersion curves → *h*, √(*E*/ρ), *ν*) |
+| `Test_data_100.mat` | 100 sample test set (500×500 dispersion curves → $h$, $\sqrt{E/\rho}$, $\nu$) |
+| `dispersion_samples.png` | Visualization of sample dispersion curves with material parameters |
 | `model.onnx` | ONNX export with embedded de-normalization |
 | `training_history.csv` | Epoch-by-epoch training metrics (loss, R², LR, etc.) |
 | `training_curves.png` | Training/validation loss and learning rate plot |

{wavedl-1.5.6 → wavedl-1.5.7}/src/wavedl/__init__.py

@@ -18,7 +18,7 @@ For inference:
     # or: python -m wavedl.test --checkpoint best_checkpoint --data_path test.npz
 """
 
-__version__ = "1.5.6"
+__version__ = "1.5.7"
 __author__ = "Ductho Le"
 __email__ = "ductho.le@outlook.com"
 

{wavedl-1.5.6 → wavedl-1.5.7}/src/wavedl/models/efficientnet.py

@@ -110,9 +110,30 @@ class EfficientNetBase(BaseModel):
             self._freeze_backbone()
 
     def _adapt_input_channels(self):
-        """Modify first conv to handle single-channel input by expanding to 3ch."""
-        # We'll handle this in forward by repeating channels
-        pass
+        """Modify first conv to accept single-channel input.
+
+        Instead of expanding 1→3 channels in forward (which triples memory),
+        we replace the first conv layer with a 1-channel version and initialize
+        weights as the mean of the pretrained RGB filters.
+        """
+        # EfficientNet stem conv is at: features[0][0]
+        old_conv = self.backbone.features[0][0]
+        new_conv = nn.Conv2d(
+            1,  # Single channel input
+            old_conv.out_channels,
+            kernel_size=old_conv.kernel_size,
+            stride=old_conv.stride,
+            padding=old_conv.padding,
+            dilation=old_conv.dilation,
+            groups=old_conv.groups,
+            padding_mode=old_conv.padding_mode,
+            bias=old_conv.bias is not None,
+        )
+        if self.pretrained:
+            # Initialize with mean of pretrained RGB weights
+            with torch.no_grad():
+                new_conv.weight.copy_(old_conv.weight.mean(dim=1, keepdim=True))
+        self.backbone.features[0][0] = new_conv
 
     def _freeze_backbone(self):
         """Freeze all backbone parameters except the classifier."""
@@ -130,10 +151,6 @@ class EfficientNetBase(BaseModel):
         Returns:
             Output tensor of shape (B, out_size)
         """
-        # Expand single channel to 3 channels for pretrained weights
-        if x.size(1) == 1:
-            x = x.expand(-1, 3, -1, -1)
-
         return self.backbone(x)
 
     @classmethod

{wavedl-1.5.6 → wavedl-1.5.7}/src/wavedl/models/efficientnetv2.py

@@ -129,10 +129,37 @@ class EfficientNetV2Base(BaseModel):
             nn.Linear(regression_hidden // 2, out_size),
         )
 
-        # Optionally freeze backbone for fine-tuning
+        # Adapt first conv for single-channel input (3× memory savings vs expand)
+        self._adapt_input_channels()
+
+        # Optionally freeze backbone for fine-tuning (after adaptation so new conv is frozen too)
         if freeze_backbone:
             self._freeze_backbone()
 
+    def _adapt_input_channels(self):
+        """Modify first conv to accept single-channel input.
+
+        Instead of expanding 1→3 channels in forward (which triples memory),
+        we replace the first conv layer with a 1-channel version and initialize
+        weights as the mean of the pretrained RGB filters.
+        """
+        old_conv = self.backbone.features[0][0]
+        new_conv = nn.Conv2d(
+            1,  # Single channel input
+            old_conv.out_channels,
+            kernel_size=old_conv.kernel_size,
+            stride=old_conv.stride,
+            padding=old_conv.padding,
+            dilation=old_conv.dilation,
+            groups=old_conv.groups,
+            padding_mode=old_conv.padding_mode,
+            bias=old_conv.bias is not None,
+        )
+        if self.pretrained:
+            with torch.no_grad():
+                new_conv.weight.copy_(old_conv.weight.mean(dim=1, keepdim=True))
+        self.backbone.features[0][0] = new_conv
+
     def _freeze_backbone(self):
         """Freeze all backbone parameters except the classifier."""
         for name, param in self.backbone.named_parameters():
@@ -144,15 +171,11 @@ class EfficientNetV2Base(BaseModel):
         Forward pass.
 
         Args:
-            x: Input tensor of shape (B, C, H, W) where C is 1 or 3
+            x: Input tensor of shape (B, 1, H, W)
 
         Returns:
             Output tensor of shape (B, out_size)
         """
-        # Expand single channel to 3 channels for pretrained weights compatibility
-        if x.size(1) == 1:
-            x = x.expand(-1, 3, -1, -1)
-
         return self.backbone(x)
 
     @classmethod

{wavedl-1.5.6 → wavedl-1.5.7}/src/wavedl/models/mobilenetv3.py

@@ -136,10 +136,37 @@ class MobileNetV3Base(BaseModel):
             nn.Linear(regression_hidden, out_size),
         )
 
-        # Optionally freeze backbone for fine-tuning
+        # Adapt first conv for single-channel input (3× memory savings vs expand)
+        self._adapt_input_channels()
+
+        # Optionally freeze backbone for fine-tuning (after adaptation so new conv is frozen too)
         if freeze_backbone:
             self._freeze_backbone()
 
+    def _adapt_input_channels(self):
+        """Modify first conv to accept single-channel input.
+
+        Instead of expanding 1→3 channels in forward (which triples memory),
+        we replace the first conv layer with a 1-channel version and initialize
+        weights as the mean of the pretrained RGB filters.
+        """
+        old_conv = self.backbone.features[0][0]
+        new_conv = nn.Conv2d(
+            1,  # Single channel input
+            old_conv.out_channels,
+            kernel_size=old_conv.kernel_size,
+            stride=old_conv.stride,
+            padding=old_conv.padding,
+            dilation=old_conv.dilation,
+            groups=old_conv.groups,
+            padding_mode=old_conv.padding_mode,
+            bias=old_conv.bias is not None,
+        )
+        if self.pretrained:
+            with torch.no_grad():
+                new_conv.weight.copy_(old_conv.weight.mean(dim=1, keepdim=True))
+        self.backbone.features[0][0] = new_conv
+
     def _freeze_backbone(self):
         """Freeze all backbone parameters except the classifier."""
         for name, param in self.backbone.named_parameters():
@@ -151,15 +178,11 @@ class MobileNetV3Base(BaseModel):
         Forward pass.
 
         Args:
-            x: Input tensor of shape (B, C, H, W) where C is 1 or 3
+            x: Input tensor of shape (B, 1, H, W)
 
         Returns:
             Output tensor of shape (B, out_size)
         """
-        # Expand single channel to 3 channels for pretrained weights compatibility
-        if x.size(1) == 1:
-            x = x.expand(-1, 3, -1, -1)
-
         return self.backbone(x)
 
     @classmethod
@@ -194,7 +217,7 @@ class MobileNetV3Small(MobileNetV3Base):
 
     Performance (approximate):
         - CPU inference: ~6ms (single core)
-        - Parameters: 2.5M
+        - Parameters: ~1.1M
         - MAdds: 56M
 
     Args:
@@ -241,7 +264,7 @@ class MobileNetV3Large(MobileNetV3Base):
 
     Performance (approximate):
         - CPU inference: ~20ms (single core)
-        - Parameters: 5.4M
+        - Parameters: ~3.2M
         - MAdds: 219M
 
     Args:

{wavedl-1.5.6 → wavedl-1.5.7}/src/wavedl/models/regnet.py

@@ -140,10 +140,37 @@ class RegNetBase(BaseModel):
             nn.Linear(regression_hidden, out_size),
         )
 
-        # Optionally freeze backbone for fine-tuning
+        # Adapt first conv for single-channel input (3× memory savings vs expand)
+        self._adapt_input_channels()
+
+        # Optionally freeze backbone for fine-tuning (after adaptation so new conv is frozen too)
         if freeze_backbone:
             self._freeze_backbone()
 
+    def _adapt_input_channels(self):
+        """Modify first conv to accept single-channel input.
+
+        Instead of expanding 1→3 channels in forward (which triples memory),
+        we replace the first conv layer with a 1-channel version and initialize
+        weights as the mean of the pretrained RGB filters.
+        """
+        old_conv = self.backbone.stem[0]
+        new_conv = nn.Conv2d(
+            1,  # Single channel input
+            old_conv.out_channels,
+            kernel_size=old_conv.kernel_size,
+            stride=old_conv.stride,
+            padding=old_conv.padding,
+            dilation=old_conv.dilation,
+            groups=old_conv.groups,
+            padding_mode=old_conv.padding_mode,
+            bias=old_conv.bias is not None,
+        )
+        if self.pretrained:
+            with torch.no_grad():
+                new_conv.weight.copy_(old_conv.weight.mean(dim=1, keepdim=True))
+        self.backbone.stem[0] = new_conv
+
     def _freeze_backbone(self):
         """Freeze all backbone parameters except the fc layer."""
         for name, param in self.backbone.named_parameters():
@@ -155,15 +182,11 @@ class RegNetBase(BaseModel):
         Forward pass.
 
         Args:
-            x: Input tensor of shape (B, C, H, W) where C is 1 or 3
+            x: Input tensor of shape (B, 1, H, W)
 
         Returns:
             Output tensor of shape (B, out_size)
         """
-        # Expand single channel to 3 channels for pretrained weights compatibility
-        if x.size(1) == 1:
-            x = x.expand(-1, 3, -1, -1)
-
         return self.backbone(x)
 
     @classmethod

{wavedl-1.5.6 → wavedl-1.5.7}/src/wavedl/models/swin.py

@@ -141,10 +141,46 @@ class SwinTransformerBase(BaseModel):
             nn.Linear(regression_hidden // 2, out_size),
         )
 
-        # Optionally freeze backbone for fine-tuning
+        # Adapt patch embedding conv for single-channel input (3× memory savings vs expand)
+        self._adapt_input_channels()
+
+        # Optionally freeze backbone for fine-tuning (after adaptation so new conv is frozen too)
         if freeze_backbone:
             self._freeze_backbone()
 
+    def _adapt_input_channels(self):
+        """Modify patch embedding conv to accept single-channel input.
+
+        Instead of expanding 1→3 channels in forward (which triples memory),
+        we replace the patch embedding conv with a 1-channel version and
+        initialize weights as the mean of the pretrained RGB filters.
+        """
+        # Swin's patch embedding is at features[0][0]
+        try:
+            old_conv = self.backbone.features[0][0]
+        except (IndexError, AttributeError, TypeError) as e:
+            raise RuntimeError(
+                f"Swin patch embed structure changed in this torchvision version. "
+                f"Cannot adapt input channels. Error: {e}"
+            ) from e
+        new_conv = nn.Conv2d(
+            1,  # Single channel input
+            old_conv.out_channels,
+            kernel_size=old_conv.kernel_size,
+            stride=old_conv.stride,
+            padding=old_conv.padding,
+            dilation=old_conv.dilation,
+            groups=old_conv.groups,
+            padding_mode=old_conv.padding_mode,
+            bias=old_conv.bias is not None,
+        )
+        if self.pretrained:
+            with torch.no_grad():
+                new_conv.weight.copy_(old_conv.weight.mean(dim=1, keepdim=True))
+                if old_conv.bias is not None:
+                    new_conv.bias.copy_(old_conv.bias)
+        self.backbone.features[0][0] = new_conv
+
     def _freeze_backbone(self):
         """Freeze all backbone parameters except the head."""
         for name, param in self.backbone.named_parameters():
@@ -156,15 +192,11 @@ class SwinTransformerBase(BaseModel):
         Forward pass.
 
         Args:
-            x: Input tensor of shape (B, C, H, W) where C is 1 or 3
+            x: Input tensor of shape (B, 1, H, W)
 
         Returns:
             Output tensor of shape (B, out_size)
         """
-        # Expand single channel to 3 channels for pretrained weights compatibility
-        if x.size(1) == 1:
-            x = x.expand(-1, 3, -1, -1)
-
         return self.backbone(x)
 
     @classmethod

{wavedl-1.5.6 → wavedl-1.5.7}/src/wavedl/models/tcn.py

@@ -45,6 +45,26 @@ from wavedl.models.base import BaseModel
 from wavedl.models.registry import register_model
 
 
+def _find_group_count(channels: int, max_groups: int = 8) -> int:
+    """
+    Find largest valid group count for GroupNorm.
+
+    GroupNorm requires channels to be divisible by num_groups.
+    This finds the largest divisor up to max_groups.
+
+    Args:
+        channels: Number of channels
+        max_groups: Maximum group count to consider (default: 8)
+
+    Returns:
+        Largest valid group count (always >= 1)
+    """
+    for g in range(min(max_groups, channels), 0, -1):
+        if channels % g == 0:
+            return g
+    return 1
+
+
 class CausalConv1d(nn.Module):
     """
     Causal 1D convolution with dilation.
@@ -101,13 +121,13 @@ class TemporalBlock(nn.Module):
 
         # First causal convolution
         self.conv1 = CausalConv1d(in_channels, out_channels, kernel_size, dilation)
-        self.norm1 = nn.GroupNorm(min(8, out_channels), out_channels)
+        self.norm1 = nn.GroupNorm(_find_group_count(out_channels), out_channels)
         self.act1 = nn.GELU()
         self.dropout1 = nn.Dropout(dropout)
 
         # Second causal convolution
         self.conv2 = CausalConv1d(out_channels, out_channels, kernel_size, dilation)
-        self.norm2 = nn.GroupNorm(min(8, out_channels), out_channels)
+        self.norm2 = nn.GroupNorm(_find_group_count(out_channels), out_channels)
         self.act2 = nn.GELU()
         self.dropout2 = nn.Dropout(dropout)
 

{wavedl-1.5.6 → wavedl-1.5.7}/src/wavedl/test.py

@@ -311,7 +311,7 @@ def load_data_for_inference(
 # ==============================================================================
 def load_checkpoint(
     checkpoint_dir: str,
-    in_shape: tuple[int, int],
+    in_shape: tuple[int, ...],
     out_size: int,
     model_name: str | None = None,
 ) -> tuple[nn.Module, any]:
@@ -320,7 +320,7 @@ def load_checkpoint(
 
     Args:
         checkpoint_dir: Path to checkpoint directory
-        in_shape: Input image shape (H, W)
+        in_shape: Input spatial shape - (L,) for 1D, (H, W) for 2D, or (D, H, W) for 3D
         out_size: Number of output parameters
         model_name: Model architecture name (auto-detect if None)
 
@@ -376,7 +376,11 @@ def load_checkpoint(
                 )
 
     logging.info(f"   Building model: {model_name}")
-    model = build_model(model_name, in_shape=in_shape, out_size=out_size)
+    # Use pretrained=False: checkpoint weights will overwrite any pretrained weights,
+    # so downloading ImageNet weights is wasteful and breaks offline/HPC inference.
+    model = build_model(
+        model_name, in_shape=in_shape, out_size=out_size, pretrained=False
+    )
 
     # Load weights (check multiple formats in order of preference)
     weight_path = None

{wavedl-1.5.6 → wavedl-1.5.7}/src/wavedl/train.py

@@ -240,15 +240,9 @@ def parse_args() -> argparse.Namespace:
     )
     parser.add_argument(
         "--pretrained",
-        action="store_true",
+        action=argparse.BooleanOptionalAction,
         default=True,
-        help="Use pretrained weights (default: True)",
-    )
-    parser.add_argument(
-        "--no_pretrained",
-        dest="pretrained",
-        action="store_false",
-        help="Train from scratch without pretrained weights",
+        help="Use pretrained weights (default: True). Use --no-pretrained to train from scratch.",
     )
 
     # Configuration File
@@ -367,6 +361,18 @@ def parse_args() -> argparse.Namespace:
         help="DataLoader workers per GPU (-1=auto-detect based on CPU cores)",
     )
     parser.add_argument("--seed", type=int, default=2025, help="Random seed")
+    parser.add_argument(
+        "--deterministic",
+        action="store_true",
+        help="Enable deterministic mode for reproducibility (slower, disables TF32/cuDNN benchmark)",
+    )
+    parser.add_argument(
+        "--cache_validate",
+        type=str,
+        default="sha256",
+        choices=["sha256", "fast", "size"],
+        help="Cache validation mode: sha256 (full hash), fast (partial), size (quick)",
+    )
     parser.add_argument(
         "--single_channel",
         action="store_true",
@@ -512,11 +518,23 @@ def main():
                     # Import as regular module
                     importlib.import_module(module_name)
                     print(f"✓ Imported module: {module_name}")
-            except ImportError as e:
+            except (ImportError, FileNotFoundError, SyntaxError, PermissionError) as e:
                 print(f"✗ Failed to import '{module_name}': {e}", file=sys.stderr)
-                print(
-                    "  Make sure the module is in your Python path or current directory."
-                )
+                if isinstance(e, FileNotFoundError):
+                    print("  File does not exist. Check the path.", file=sys.stderr)
+                elif isinstance(e, SyntaxError):
+                    print(
+                        f"  Syntax error at line {e.lineno}: {e.msg}", file=sys.stderr
+                    )
+                elif isinstance(e, PermissionError):
+                    print(
+                        "  Permission denied. Check file permissions.", file=sys.stderr
+                    )
+                else:
+                    print(
+                        "  Make sure the module is in your Python path or current directory.",
+                        file=sys.stderr,
+                    )
                 sys.exit(1)
 
     # Handle --list_models flag
@@ -648,6 +666,17 @@ def main():
     )
     set_seed(args.seed)
 
+    # Deterministic mode for scientific reproducibility
+    # Disables TF32 and cuDNN benchmark for exact reproducibility (slower)
+    if args.deterministic:
+        torch.backends.cudnn.benchmark = False
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cuda.matmul.allow_tf32 = False
+        torch.backends.cudnn.allow_tf32 = False
+        torch.use_deterministic_algorithms(True, warn_only=True)
+        if accelerator.is_main_process:
+            print("🔒 Deterministic mode enabled (slower but reproducible)")
+
     # Configure logging (rank 0 only prints to console)
     logging.basicConfig(
         level=logging.INFO if accelerator.is_main_process else logging.ERROR,

{wavedl-1.5.6 → wavedl-1.5.7}/src/wavedl/utils/constraints.py

@@ -207,22 +207,28 @@ class ExpressionConstraint(nn.Module):
             # Parse indices from the slice
             indices = self._parse_subscript_indices(node.slice)
 
+            # Auto-squeeze channel dimension for single-channel inputs
+            # This allows x[i,j] syntax for (B, 1, H, W) inputs instead of x[c,i,j]
+            inputs_for_indexing = inputs
+            if inputs.ndim >= 3 and inputs.shape[1] == 1:
+                inputs_for_indexing = inputs.squeeze(1)  # (B, 1, H, W) → (B, H, W)
+
             # Validate dimensions match
             # inputs shape: (batch, dim1) or (batch, dim1, dim2) or (batch, dim1, dim2, dim3)
-            input_ndim = inputs.ndim - 1  # Exclude batch dimension
+            input_ndim = inputs_for_indexing.ndim - 1  # Exclude batch dimension
             if len(indices) != input_ndim:
                 raise ValueError(
-                    f"Input has {input_ndim}D shape, but got {len(indices)} indices. "
+                    f"Input has {input_ndim}D shape (after channel squeeze), but got {len(indices)} indices. "
                     f"Use x[i] for 1D, x[i,j] for 2D, x[i,j,k] for 3D inputs."
                 )
 
             # Extract the value at the specified indices (for entire batch)
             if len(indices) == 1:
-                return inputs[:, indices[0]]
+                return inputs_for_indexing[:, indices[0]]
             elif len(indices) == 2:
-                return inputs[:, indices[0], indices[1]]
+                return inputs_for_indexing[:, indices[0], indices[1]]
             elif len(indices) == 3:
-                return inputs[:, indices[0], indices[1], indices[2]]
+                return inputs_for_indexing[:, indices[0], indices[1], indices[2]]
             else:
                 raise ValueError("Only 1D, 2D, or 3D input indexing supported.")
         elif isinstance(node, ast.Expression):