Dist ckpt refactor (#8)

* refactored for distributed checkpointing

* addressing comments

* missed commit

Author: adil-a

vectorlm/utils/__init__.py

@@ -12,12 +12,10 @@
     checkpoint_exists,
     get_latest_checkpoint_dir,
     load_metadata,
-    load_model,
-    load_optimizer,
+    load_model_and_optimizer,
     load_scheduler,
     save_consolidated_model,
     save_metadata,
-    save_model,
-    save_optimizer,
+    save_model_and_optimizer,
     save_scheduler,
 )