diff --git a/sdm/models/tabfm/row_embedding.py b/sdm/models/tabfm/row_embedding.py
index d49a7d22d4..3ae887f491 100644
--- a/sdm/models/tabfm/row_embedding.py
+++ b/sdm/models/tabfm/row_embedding.py
@@ -27,14 +27,13 @@
     Linear,
     ModuleList,
     Parameter,
-    RMSNorm,
     Sequential,
 )
 
 from sdm.cache import Cache, KVCacheEntry
 from sdm.models.tabfm.block import TabFMTransformerBlock
 from sdm.models.tabfm.cell_embedding import CellEmbedding
-from sdm.nn import InducedTransformerBlock, RotaryEmbedding
+from sdm.nn import InducedTransformerBlock, RMSNorm, RotaryEmbedding
 
 
 class RowEmbedding(torch.nn.Module):