Fix bf16 standalone RMSNorm precision

README.md

@@ -17,6 +17,7 @@ library_name: fcdm_diffae
 
 | Date | Change |
 |------|--------|
+| 2026-04-10 | Refresh standalone package: fix bf16 RMSNorm precision path in both encoder and decoder to match training code; local export tooling now preserves fp32 EMA weights for future re-exports |
 | 2026-04-08 | Initial release |
 
 **Experimental patch-32 version** of

fcdm_diffae/norms.py

@@ -30,10 +30,10 @@ class ChannelWiseRMSNorm(nn.Module):
         # Float32 accumulation for stability
         ms = torch.mean(torch.square(x), dim=1, keepdim=True, dtype=torch.float32)
         inv_rms = torch.rsqrt(ms + self._eps)
-        y = x * inv_rms
+        y = x * inv_rms.to(dtype=x.dtype)
         if self.weight is not None:
             shape = (1, -1) + (1,) * (x.dim() - 2)
-            y = y * self.weight.view(shape).to(dtype=y.dtype)
+            y = y * self.weight.view(shape).to(dtype=x.dtype)
             if self.bias is not None:
-                y = y + self.bias.view(shape).to(dtype=y.dtype)
-        return y.to(dtype=x.dtype)
+                y = y + self.bias.view(shape).to(dtype=x.dtype)
+        return y