from tinygrad.nn.state import safe_save
import csv
import os

def logLoss(step, loss):
    path = "loss.csv"
    exists = os.path.isfile(path)
    with open(path, mode='a', newline='') as f:
        writer = csv.writer(f)
        if not exists:
            writer.writerow(['step', 'loss'])
        writer.writerow([step, float(loss)])

def logModel(step,stateDict):
    safe_save(stateDict, f"gpt_{step}.safetensors")