cleanup debug code, v nice state dict hook workaround

Author: ebsmothers

torchtune/modules/model_fusion/_fusion.py

@@ -232,10 +232,11 @@ def _load_state_dict_hook(self, state_dict, prefix, *args, **kwargs):
         """Apply extra "embedding" prefix to the state_dict key to
         account for the FusionEmbedding wrapping.
         """
-        key = prefix + "weight"
-        new_key = prefix + "embedding.weight"
-        state_dict[new_key] = state_dict[key]
-        del state_dict[key]
+        if state_dict:
+            key = prefix + "weight"
+            new_key = prefix + "embedding.weight"
+            state_dict[new_key] = state_dict[key]
+            del state_dict[key]
 
     def fusion_params(self) -> List[str]:
         """

torchtune/modules/peft/_utils.py

@@ -259,7 +259,6 @@ def get_merged_lora_ckpt(
 
         # Otherwise it is just vanilla LoRA
         else:
-            print(f"module is {module}")
             state_dict[f"{module}.weight"] += (
                 (alpha / rank) * lora_b_weight @ lora_a_weight
             )