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

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