From 9e623f2ff680b1cc6422c89b644100dd6ffb8b4c Mon Sep 17 00:00:00 2001 From: yesMohanHere Date: Sat, 18 Jul 2026 00:09:21 +0530 Subject: [PATCH] Fix custom_bwd decorator usage across torch versions torch.cuda.amp.custom_bwd (torch < 2.4) is applied bare, while torch.amp.custom_bwd (torch >= 2.4) must be called with device_type to yield the decorator. Use small wrappers so both work. --- nerfstudio/field_components/activations.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/nerfstudio/field_components/activations.py b/nerfstudio/field_components/activations.py index a7a54c319..b8edebc7f 100644 --- a/nerfstudio/field_components/activations.py +++ b/nerfstudio/field_components/activations.py @@ -25,27 +25,28 @@ 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 + from torch.amp import custom_bwd as _custom_bwd, custom_fwd as _custom_fwd - _CUSTOM_AMP_KWARGS = {"device_type": "cuda"} + def custom_fwd(*args, **kwargs): + return _custom_fwd(*args, device_type="cuda", **kwargs) + + custom_bwd = _custom_bwd(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, **_CUSTOM_AMP_KWARGS) + @custom_fwd(cast_inputs=torch.float32) def forward(ctx, x): ctx.save_for_backward(x) return torch.exp(x) @staticmethod - @custom_bwd(**_CUSTOM_AMP_KWARGS) + @custom_bwd def backward(ctx, g): x = ctx.saved_tensors[0] return g * torch.exp(x.clamp(-15, 15))