Added error handling on run.py

run.py

@@ -22,6 +22,7 @@ CKPT_PATH = "./checkpoints/"
 
 
 def main():
+    try:
         grok_1_model = LanguageModelConfig(
             vocab_size=128 * 1024,
             pad_token=0,
@@ -47,6 +48,7 @@ def main():
                 model_axis="model",
             ),
         )
+
         inference_runner = InferenceRunner(
             pad_sizes=(1024,),
             runner=ModelRunner(
@@ -60,12 +62,21 @@ def main():
             local_mesh_config=(1, 8),
             between_hosts_config=(1, 1),
         )
+
         inference_runner.initialize()
         gen = inference_runner.run()
 
         inp = "The answer to life the universe and everything is of course"
         print(f"Output for prompt: {inp}", sample_from_model(gen, inp, max_len=100, temperature=0.01))
 
+    except FileNotFoundError as e:
+        logging.error(f"File not found: {e}")
+    except ValueError as e:
+        logging.error(f"Value error: {e}")
+    except Exception as e:
+        logging.error(f"An unexpected error occurred: {e}")
+
+
 
 if __name__ == "__main__":
     logging.basicConfig(level=logging.INFO)