Update modeling_pragya.py

modeling_pragya.py

@@ -68,8 +68,12 @@ class PragyaForCausalLM(PreTrainedModel, GenerationMixin):
             for name, buf in _non_persistent_buffers(fresh).items():
                 own[name].copy_(buf)
         del fresh
-        # wte and lm_head are one tensor (checkpoint stores only lm_head). Re-tie after load.
-        self.model.transformer.wte.weight = self.model.lm_head.weight
+        # wte and lm_head are one tensor. Checkpoints store one of the two, so tie from the
+        # one that was actually loaded, not from whichever happens to be first.
+        if "model.lm_head.weight" in missing_keys:
+            self.model.lm_head.weight = self.model.transformer.wte.weight
+        else:
+            self.model.transformer.wte.weight = self.model.lm_head.weight
 
     def _init_weights(self, module):
         # Weights come from GPT.__init__ (Engram tables, zero-init attention gate, tied