mup base shapes path

modeling_nt_bert.py

@@ -78,11 +78,13 @@ class BertPreTrainedModel(PreTrainedModel):
 
         # since we used MuP, need to reset values since they're not saved with the model
         if os.path.exists("base_shapes.bsh") is False:
-            hf_hub_download(
+            path = hf_hub_download(
                 "zpn/human_bp_bert", "base_shapes.bsh"
             )
+        else:
+           path = "base_shapes.bsh"
             
-        set_base_shapes(model, "base_shapes.bsh", rescale_params=False)
+        set_base_shapes(model, path, rescale_params=False)
 
         return model