Skip to content

[Bug] JukeboxLoss.forward: input_amplitude and target_amplitude variable names are swapped #8820

@Zeesejo

Description

@Zeesejo

Describe the bug

In monai/losses/spectral_loss.py, JukeboxLoss.forward() has the variable names input_amplitude and target_amplitude swapped:

def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    input_amplitude = self._get_fft_amplitude(target)   # ← BUG: should be _get_fft_amplitude(input)
    target_amplitude = self._get_fft_amplitude(input)   # ← BUG: should be _get_fft_amplitude(target)

    loss = F.mse_loss(target_amplitude, input_amplitude, reduction="none")

Impact

While this does not affect the numerical output of the loss (since MSE is symmetric: mse(a, b) == mse(b, a)), the naming is semantically inverted which:

  1. Makes the code misleading and confusing to anyone reading, extending, or debugging it
  2. Creates a silent maintenance hazard — if anyone adds asymmetric processing (e.g. normalization, weighting, gradient masking) based on input vs. target, the swap will cause subtle, hard-to-debug training errors
  3. Violates the principle of least surprise for the standard forward(input, target) contract that all PyTorch loss functions follow

Root Cause

In forward(self, input, target):

  • input_amplitude is computed from target (wrong)
  • target_amplitude is computed from input (wrong)

Expected behavior

def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    input_amplitude = self._get_fft_amplitude(input)    # correct
    target_amplitude = self._get_fft_amplitude(target)  # correct
    loss = F.mse_loss(target_amplitude, input_amplitude, reduction="none")

Location

monai/losses/spectral_loss.py lines 55-56

Affects all MONAI versions containing JukeboxLoss. Reproducible on dev branch as of 2026-04-11.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions