Minor refactor

llmc_py/utils.py

@@ -1,3 +1,7 @@
+# Taken from:
+# 1) https://github.com/meta-llama/llama-models/blob/main/models/llama3_1/api/model.py
+# 2) https://github.com/meta-llama/llama3/blob/11817d47e1ba7a4959b025eb1ca308572e0e3963/llama/generation.py
+
 import torch
 from torch import nn
 

train_gpt2.py

@@ -24,7 +24,6 @@
 from contextlib import nullcontext
 from dataclasses import dataclass
 from pathlib import Path
-
 from typing import (
     List,
     Optional,
@@ -988,6 +987,8 @@ def print0(*args, **kwargs):
             print(f"{result['generation']}")
             print("\n==================================\n")
 
+        exit(0)  # only inference supported for now
+
     # -------------------------------------------------------------------------
     # PyTorch -> C bridge: save some weights and state for C to load later as reference