fix(rms_norm.py): add assertion for input gradients to handle unsupported cases in backward pass

torch-ext/activation/rms_norm.py

@@ -59,6 +59,9 @@ class FusedAddRMSNormFunction(torch.autograd.Function):
         need_in = ctx.needs_input_grad[0]
         need_res = ctx.needs_input_grad[1]
 
+        # TODO(ai-system): kernels currently do not support no input gradients
+        assert need_in or need_res, "Not implemented for no input gradients yet"
+
         grad, weight_grad = ops.fused_add_rms_norm_backward(
             output_grad,
             add_output_grad,