cleanup

Signed-off-by: Mayank Mishra <[email protected]>

dolomite_engine/hf_models/mixins/dense_TP/base.py

@@ -37,7 +37,7 @@ def __init__(self, config: CommonConfig, *args, **kwargs):
 
         super().__init__(config, *args, **kwargs)
 
-        if self._tied_word_embeddings:
+        if self.pp_world_size > 1 and self._tied_word_embeddings:
             raise NotImplementedError()