Replace deprecated torch.cuda.amp.GradScaler with torch.amp in capture.py by xyf5432 · Pull Request #1946 · NVIDIA/physicsnemo

xyf5432 · 2026-08-21T08:55:37Z

Fixes #1945

Summary

torch.cuda.amp has been deprecated since torch 2.3 (GradScaler) / 2.4 (autocast) and is scheduled for removal. This PR migrates every remaining use in the repository — the core library, 13 example scripts, 2 Hydra configs, and the test suite — to the torch.amp API:

Core library — physicsnemo/utils/capture.py:272-274, the last core reference:

-    ) -> torch.cuda.amp.GradScaler:
+    ) -> torch.amp.GradScaler:
         # Create gradient scaler
-        scaler = torch.cuda.amp.GradScaler(enabled=scaler_enabled)
+        scaler = torch.amp.GradScaler("cuda", enabled=scaler_enabled)

Examples (13 files under examples/): from torch.cuda.amp import GradScaler/autocast → from torch.amp import ... (bare call sites unchanged), the partial(torch.cuda.amp.autocast, ...) in figconvnet/train.py → partial(torch.amp.autocast, "cuda", ...), and type annotations/docstrings in darcy_transolver/train_transolver_darcy_fix.py.
Hydra configs (2 files): _target_: torch.cuda.amp.GradScaler → _target_: torch.amp.GradScaler (instantiate passes enabled=...; device defaults to "cuda").
Tests — test/common/optimization.py (2 call sites).

torch.amp.GradScaler("cuda", ...) / torch.amp.autocast("cuda", ...) are the official replacements (available since torch 2.3 / 2.0, well below the torch>=2.10.0 floor in pyproject.toml) and share the same methods and state-dict/checkpoint format, so behavior is unchanged. The only remaining torch.cuda.amp string in the repo is a doc URL comment in examples/minimal/ShardTensorExamples/5_vit_training_loop/utils/measure_perf.py, not code.

Type of change

Bug fix (non-breaking change which fixes an issue)

Validation

Verified on torch 2.11.0 with warnings.simplefilter("error", FutureWarning):

torch.amp.GradScaler("cuda", enabled=...) and torch.amp.autocast("cuda", ...) construct warning-free (AMP is disabled when CUDA is unavailable, by design).
Control group: torch.cuda.amp.GradScaler(enabled=True) raises FutureWarning under the same filter — the fix is effective and the filter is sensitive.
py_compile passes on all 15 touched files; the repo now has zero code references to torch.cuda.amp (grep-verified).
ruff check shows no new errors vs. baseline; all touched files pass ruff format --check (the one remaining unformatted file, train_graphcast.py, was already unformatted at HEAD).

User impact

No behavior change; the FutureWarning emitted at import of physicsnemo.utils.capture and on every AMP-enabled example/test run on torch 2.4+ is eliminated, as is the AttributeError/Hydra instantiation failure risk once torch.cuda.amp is removed.

Notes for reviewers

Same migration as ultralytics/yolov5#13244 and huggingface/lerobot#3167.

torch.cuda.amp.GradScaler is deprecated since torch 2.3 and scheduled for removal; torch.amp.GradScaler('cuda') is the API-compatible replacement (available since torch 2.3, well below the torch>=2.10 floor). This removes the last torch.cuda.amp reference in the core physicsnemo library (examples/tests still use it, tracked in NVIDIA#1945). Fixes NVIDIA#1945 Co-Authored-By: Claude <noreply@anthropic.com>

copy-pr-bot · 2026-08-21T08:55:41Z

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.

github-actions · 2026-08-21T08:55:48Z

CODEOWNERS review map

Current for commit 42a13880ecf3. An approval covers every file listed for that owner; one owner is sufficient for shared files.

⏳ @ktangsali — 1 file(s)

physicsnemo/utils/capture.py

⏳ @NickGeneva — 1 file(s)

physicsnemo/utils/capture.py

Comment /codeowners-info to refresh.

greptile-apps · 2026-08-21T08:57:52Z

Greptile Summary

Migrates static-capture gradient scaling from the deprecated torch.cuda.amp namespace to the supported torch.amp API.

Updates the _init_amp_scaler return annotation.
Constructs the scaler with an explicit CUDA device type.
Preserves existing enabled-state and checkpoint behavior.

Important Files Changed

Filename	Overview
physicsnemo/utils/capture.py	Replaces the deprecated GradScaler type and constructor with the compatible `torch.amp` equivalents without changing capture behavior.

_{Reviews (1): Last reviewed commit: "Use torch.amp.GradScaler in capture.py" | Re-trigger Greptile}

Completes the torch.cuda.amp deprecation cleanup tracked in NVIDIA#1945: mechanical replacement across 13 example scripts, 2 Hydra configs (_target_ class paths) and test/common/optimization.py. The remaining torch.cuda.amp string is a doc URL comment in examples/minimal/ShardTensorExamples/.../measure_perf.py, not code. Co-Authored-By: Claude <noreply@anthropic.com>

peterdsharpe · 2026-08-21T13:09:36Z

/ok to test 5a5458e

peterdsharpe

LGTM, thanks

peterdsharpe · 2026-08-25T21:16:13Z

/ok to test de9a8f4

peterdsharpe · 2026-09-01T21:55:10Z

/ok to test 0aee19a

Copilot

🟡 Changes recommended

Several scripts now import torch.amp.autocast but still call autocast(...) without providing/binding device_type, which will error at runtime.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

This PR migrates remaining uses of deprecated torch.cuda.amp APIs to torch.amp across the core library, examples, Hydra configs, and tests to eliminate FutureWarnings and avoid future breakage when torch.cuda.amp is removed.

Changes:

Update core AMP scaler creation/typing in physicsnemo/utils/capture.py to torch.amp.GradScaler.
Replace torch.cuda.amp imports/usages in multiple example training scripts and one benchmark utility.
Update Hydra _target_ references and adjust test AMP scaler construction.

File summaries

File	Description
`physicsnemo/utils/capture.py`	Switch core gradient-scaler annotation/instantiation to `torch.amp.GradScaler`.
`test/common/optimization.py`	Update test scaler construction away from deprecated `torch.cuda.amp.GradScaler`.
`examples/weather/graphcast/train_graphcast.py`	Update GradScaler import to `torch.amp`.
`examples/weather/graphcast/train_base.py`	Update autocast import to `torch.amp`.
`examples/reservoir_simulation/xmgn/src/train.py`	Update GradScaler import to `torch.amp`.
`examples/healthcare/bloodflow_1d_mgn/train.py`	Update GradScaler import to `torch.amp`.
`examples/healthcare/bloodflow_1d_mgn/inference.py`	Update GradScaler import to `torch.amp`.
`examples/cfd/transient_conjugate_heat_transfer_tank_fill/train.py`	Update GradScaler/autocast imports to `torch.amp`.
`examples/cfd/external_aerodynamics/xaeronet/volume/train.py`	Update GradScaler import to `torch.amp`.
`examples/cfd/external_aerodynamics/moe/train.py`	Update GradScaler/autocast imports to `torch.amp`.
`examples/cfd/external_aerodynamics/figconvnet/train.py`	Update autocast partial to use `torch.amp.autocast` with CUDA device type.
`examples/cfd/external_aerodynamics/figconvnet/configs/base.yaml`	Update Hydra target to `torch.amp.GradScaler`.
`examples/cfd/external_aerodynamics/domino/src/conf/config.yaml`	Update Hydra target to `torch.amp.GradScaler`.
`examples/cfd/external_aerodynamics/domino/src/benchmark_dataloader.py`	Update GradScaler/autocast imports to `torch.amp`.
`examples/cfd/external_aerodynamics/domino_nim_finetuning/src/train.py`	Update GradScaler/autocast imports to `torch.amp`.
`examples/cfd/darcy_transolver/train_transolver_darcy_fix.py`	Update GradScaler type annotations/docstrings to `torch.amp.GradScaler`.

Review details

Files reviewed: 16/16 changed files
Comments generated: 5
Review effort level: Lite

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

 from hydra.utils import to_absolute_path
 from omegaconf import DictConfig, OmegaConf
 import torch.distributed as dist
-from torch.cuda.amp import GradScaler, autocast
+from torch.amp import GradScaler, autocast
 from torch.nn.parallel import DistributedDataParallel
 from torch.utils.data import DataLoader


 from torch.utils.data import DataLoader, DistributedSampler
 import numpy as np
 from tqdm import tqdm
-from torch.cuda.amp import GradScaler, autocast
+from torch.amp import GradScaler, autocast



 from omegaconf import DictConfig, OmegaConf
-from torch.cuda.amp import GradScaler, autocast
+from torch.amp import GradScaler, autocast
 from torch.nn.parallel import DistributedDataParallel
 from torch.utils.data.distributed import DistributedSampler
 from torch.utils.tensorboard import SummaryWriter


 import torch
 from torch.profiler import profile, record_function, ProfilerActivity
-from torch.cuda.amp import autocast
+from torch.amp import autocast

 import sys


        amp_device = "cpu"
        amp_dtype = torch.bfloat16
-        scaler = torch.cuda.amp.GradScaler(enabled=False)  # Always false on CPU
+        scaler = torch.amp.GradScaler("cuda", enabled=False)  # Always false on CPU


… consistency torch.amp.autocast requires device_type as first positional argument; torch.cuda.amp.autocast implied it. The migration left 6 call sites in example scripts without it, which raises TypeError at runtime. - domino_nim_finetuning/src/train.py:620,685 - external_aerodynamics/moe/train.py:189 - transient_conjugate_heat_transfer_tank_fill/train.py:82,150 - weather/graphcast/train_base.py:31 Also align test/common/optimization.py:308 GradScaler device with the CPU branch ("cpu", consistent with the rest of the file). Verified on torch 2.11.0+cpu: all six forms run clean; AST scan shows no autocast call without device_type.

xyf5432 · 2026-09-02T03:18:49Z

Fixed in 543b097:

Added the required device_type to all 6 torch.amp.autocast call sites flagged by Copilot: domino_nim_finetuning/src/train.py:620/685, external_aerodynamics/moe/train.py:189, transient_conjugate_heat_transfer_tank_fill/train.py:82/150, weather/graphcast/train_base.py:31.
Aligned test/common/optimization.py:308 to GradScaler("cpu", enabled=False) in the CPU branch, consistent with the rest of the file.

Verified on torch 2.11.0+cpu: all six autocast forms now run without TypeError (previously missing 1 required positional argument: 'device_type'); AST scan confirms no autocast call without device_type remains in the repo.

xyf5432 requested review from NickGeneva and ktangsali as code owners August 21, 2026 08:55

xyf5432 requested review from coreyjadams, mnabian and peterdsharpe as code owners August 21, 2026 09:07

peterdsharpe approved these changes Aug 21, 2026

View reviewed changes

peterdsharpe enabled auto-merge August 25, 2026 21:15

Merge branch 'main' into fix/amp-deprecation

de9a8f4

coreyjadams approved these changes Aug 31, 2026

View reviewed changes

Merge branch 'main' into fix/amp-deprecation

9ad9f5f

mnabian approved these changes Aug 31, 2026

View reviewed changes

ktangsali approved these changes Aug 31, 2026

View reviewed changes

Merge branch 'main' into fix/amp-deprecation

0aee19a

Copilot AI lite review requested due to automatic review settings September 1, 2026 21:54

Copilot started reviewing on behalf of peterdsharpe September 1, 2026 21:55 View session

Copilot AI reviewed Sep 1, 2026

View reviewed changes

auto-merge was automatically disabled September 2, 2026 03:13
Head branch was pushed to by a user without write access

Merge branch 'main' into fix/amp-deprecation

d9970c1

Navigation Menu

Uh oh!

Replace deprecated torch.cuda.amp.GradScaler with torch.amp in capture.py - #1946

Replace deprecated torch.cuda.amp.GradScaler with torch.amp in capture.py#1946
xyf5432 wants to merge 7 commits into
NVIDIA:mainfrom
xyf5432:fix/amp-deprecation

xyf5432 commented Aug 21, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented Aug 21, 2026

Uh oh!

github-actions Bot commented Aug 21, 2026

Uh oh!

greptile-apps Bot commented Aug 21, 2026

Uh oh!

peterdsharpe commented Aug 21, 2026

Uh oh!

peterdsharpe left a comment

Uh oh!

peterdsharpe commented Aug 25, 2026

Uh oh!

peterdsharpe commented Sep 1, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

xyf5432 commented Sep 2, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

Uh oh!

Conversation

xyf5432 commented Aug 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Type of change

Validation

User impact

Notes for reviewers

Uh oh!

copy-pr-bot Bot commented Aug 21, 2026

Uh oh!

github-actions Bot commented Aug 21, 2026

CODEOWNERS review map

Uh oh!

greptile-apps Bot commented Aug 21, 2026

Greptile Summary

Important Files Changed

Uh oh!

peterdsharpe commented Aug 21, 2026

Uh oh!

peterdsharpe left a comment

Choose a reason for hiding this comment

Uh oh!

peterdsharpe commented Aug 25, 2026

Uh oh!

peterdsharpe commented Sep 1, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Uh oh!

xyf5432 commented Sep 2, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

xyf5432 commented Aug 21, 2026 •

edited

Loading