[layers] support zero-dim rmsnorm (#16978)

This commit is contained in:
Yinghai Lu
2026-01-12 15:53:19 -08:00
committed by GitHub
parent c6a64e9f69
commit d3c08fb07c

View File

@@ -106,6 +106,8 @@ class RMSNorm(MultiPlatformOp):
residual: Optional[torch.Tensor] = None,
**kwargs,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
if x.numel() == 0:
return x
if self.variance_size_override is not None:
return self.forward_native(x, residual, **kwargs)
if is_batch_invariant_mode_enabled():