Hacktoberfest 2026: the issues maintainers tagged for October, open and beginner-friendly. Browse Hacktoberfest issues

DiffusionUNet3D cannot run under torch.autocast (amp_mode is not exposed)

Open
  • Is something wrong?
#1,935 1 comment 0 reactions 0 assignees View on GitHub

Maintainers usually reply within 2 days

Nobody has claimed this yet.

  • #1937 by @pjaysaval — closed without merging

Assessment

Difficulty
3/5
Estimated time
1-2 days
Newbie friendliness
72/100
Issue type
Bug
Clarity
Clearly specified
Activity status
Quiet
Tech stack
python, pytorch
Domain
machine-learning

Research direction

Start with physicsnemo/experimental/models/diffusion_unets/diffusion_unet_3d.py, focusing on DiffusionUNet3D and UNetBlock3D, then compare their constructors and layer setup with SongUNet and physicsnemo.nn UNetBlock. Check the amp_mode validation in physicsnemo/nn/module/fully_connected_layers.py and its utility module. Done means the 3D model exposes and propagates amp_mode, and the reported torch.autocast example no longer raises.

Written by the indexing model from the issue text.

Description

stale
Version

2.1.1 (also reproduces on current main)

On which installation method(s) does this occur?

Pip

Describe the issue

The experimental 3D diffusion U-Net (physicsnemo.experimental.models.diffusion_unets.DiffusionUNet3D) cannot be trained under torch.autocast. Its internal physicsnemo.nn layers (Linear, PositionalEmbedding, FourierEmbedding) validate their amp_mode flag in forward and raise when autocast is active, but neither DiffusionUNet3D nor UNetBlock3D exposes the flag, so it is always False and mixed-precision training always fails.

The 2D counterparts (SongUNet, physicsnemo.nn UNetBlock) expose amp_mode as a constructor argument and thread it through their layers; the 3D model is missing that plumbing.

Minimum reproducible example
import torch
from physicsnemo.experimental.models.diffusion_unets import DiffusionUNet3D

model = DiffusionUNet3D(
    x_channels=1, num_levels=2, model_channels=16, channel_mult=[1, 2],
    num_blocks=1,
).cuda()
x = torch.randn(2, 1, 16, 16, 16, device="cuda")
t = torch.rand(2, device="cuda")
with torch.autocast("cuda", dtype=torch.bfloat16):
    model(x, t)  # raises
Relevant log output
  File ".../physicsnemo/experimental/models/diffusion_unets/diffusion_unet_3d.py", line 515, in forward
    emb = silu(self.map_layer0(emb))
  File ".../physicsnemo/nn/module/fully_connected_layers.py", line 521, in forward
    _validate_amp(self.amp_mode)
  File ".../physicsnemo/nn/module/utils/utils.py", line 209, in _validate_amp
    raise RuntimeError(
RuntimeError: amp_mode=False but torch autocast is enabled on: cuda. Disable autocast for this region or set amp_mode=True if mixed precision is intended.
Environment details
nvidia-physicsnemo 2.1.1, torch 2.11.0+cu128 / 2.13.0, Python 3.12
Observed on A100 (Linux, CUDA 12.8) and on CPU with torch.autocast("cpu", dtype=torch.bfloat16).
Dominant language
Python
Stars
3.4k
Forks
824
Avg merge
4d 5h
Merged PRs (30d)
82

Getting set up

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from NVIDIA/physicsnemo

All issues in NVIDIA/physicsnemo

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.