Skip to content

Fix SBNoiseSchedule t_max accessors#15841

Open
fallintoplace wants to merge 1 commit into
NVIDIA-NeMo:mainfrom
fallintoplace:fix/sb-noise-schedule-t-max
Open

Fix SBNoiseSchedule t_max accessors#15841
fallintoplace wants to merge 1 commit into
NVIDIA-NeMo:mainfrom
fallintoplace:fix/sb-noise-schedule-t-max

Conversation

@fallintoplace

Copy link
Copy Markdown

Summary

Fix the SBNoiseSchedule t_max accessors so they no longer crash on access and so internal t_max calculations follow the caller's device and dtype.

Root cause

alpha_t_max and sigma_t_max were building t_max with alpha.device inside property scope, where alpha was undefined. That made both accessors raise NameError as soon as they were called.

What changed

  • add a shared _time_max_like() helper plus alpha_t_max_like() and sigma_t_max_like() accessors
  • keep the existing alpha_t_max and sigma_t_max properties, but route them through the shared helper
  • reuse the new helpers in alpha_bar_from_alpha() and sigma_bar_from_sigma() so t_max stays aligned with the reference tensor's device and dtype
  • add regression coverage for CPU float32/float64, with CUDA cases included and skipped automatically when CUDA is unavailable

Validation

  • uv run pytest tests/collections/audio/test_audio_parts_submodules_schroedinger_bridge.py -k t_max_accessors -v
  • uv run pytest tests/collections/audio/test_audio_parts_submodules_schroedinger_bridge.py -v
  • uv run --with black==24.10.0 black --check nemo/collections/audio/parts/submodules/schroedinger_bridge.py tests/collections/audio/test_audio_parts_submodules_schroedinger_bridge.py
  • uv run isort --check-only nemo/collections/audio/parts/submodules/schroedinger_bridge.py tests/collections/audio/test_audio_parts_submodules_schroedinger_bridge.py

@copy-pr-bot

copy-pr-bot Bot commented Jun 27, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants