From 63a6aca37da02511d540ebfb6d2c0d37ecaa4f06 Mon Sep 17 00:00:00 2001 From: yesMohanHere Date: Fri, 17 Jul 2026 21:56:00 +0530 Subject: [PATCH] Keep custom_fwd/custom_bwd compatible with torch < 2.4 torch.amp only exports custom_fwd/custom_bwd (with mandatory device_type) from torch 2.4; older torch only has them in torch.cuda.amp without device_type. Fall back to the legacy import and drop device_type there so nerfstudio keeps working with its declared torch>=1.13.1 floor (e.g. the cu118 Docker image, torch 2.1.2). --- nerfstudio/field_components/activations.py | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/nerfstudio/field_components/activations.py b/nerfstudio/field_components/activations.py index 371999094..a7a54c319 100644 --- a/nerfstudio/field_components/activations.py +++ b/nerfstudio/field_components/activations.py @@ -21,21 +21,31 @@ from typing import TYPE_CHECKING import torch from jaxtyping import Float from torch import Tensor -from torch.amp import custom_bwd, custom_fwd from torch.autograd import Function +try: + # torch >= 2.4: custom_fwd/custom_bwd live in torch.amp and require device_type + from torch.amp import custom_bwd, custom_fwd + + _CUSTOM_AMP_KWARGS = {"device_type": "cuda"} +except ImportError: + # torch < 2.4: custom_fwd/custom_bwd live in torch.cuda.amp (no device_type argument) + from torch.cuda.amp import custom_bwd, custom_fwd + + _CUSTOM_AMP_KWARGS = {} + class _TruncExp(Function): # Implementation from torch-ngp: # https://github.com/ashawkey/torch-ngp/blob/93b08a0d4ec1cc6e69d85df7f0acdfb99603b628/activation.py @staticmethod - @custom_fwd(cast_inputs=torch.float32, device_type="cuda") + @custom_fwd(cast_inputs=torch.float32, **_CUSTOM_AMP_KWARGS) def forward(ctx, x): ctx.save_for_backward(x) return torch.exp(x) @staticmethod - @custom_bwd(device_type="cuda") + @custom_bwd(**_CUSTOM_AMP_KWARGS) def backward(ctx, g): x = ctx.saved_tensors[0] return g * torch.exp(x.clamp(-15, 15))