Update modeling_rcps.py

modeling_rcps.py

@@ -148,6 +148,11 @@ class RCPSMambaBlock(nn.Module):
         self.mixer = RCPSWrapper(mixer_cls(dim))
         norm_f = norm_cls(dim)
         self.norm = norm_f if fused_add_norm else RCPSAddNormWrapper(norm_f)
+        if self.fused_add_norm:
+            assert RMSNorm is not None, "RMSNorm import fails"
+            assert isinstance(
+                self.norm, (nn.LayerNorm, RMSNorm)
+            ), "Only LayerNorm and RMSNorm are supported for fused_add_norm"
 
     def forward(
         self, hidden_states: Tensor, residual: Optional[Tensor] = None, inference_params=None