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

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).

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.