Removed dotenv dependency

pyproject.toml

@@ -15,8 +15,7 @@ dependencies = [
     "ipykernel",
     "jaxtyping",
     "tiktoken",
-    "transformers",
-    "python-dotenv"
+    "transformers"
 ]
 
 [project.optional-dependencies]

simple_stories_train/train_llama.py

@@ -63,7 +63,7 @@
 from torch.nn import functional as F
 from torch.nn.parallel import DistributedDataParallel as DDP
 
-from simple_stories_train.utils import (
+from utils import (
     init_wandb,
     is_checkpoint_step,
     log_generations,
@@ -1036,7 +1036,7 @@ def get_lr(it: int) -> float:
                 # addition of gradients corresponds to a SUM in the objective, but
                 # instead of a SUM we want MEAN, so we scale the loss here
                 loss = loss / grad_accum_steps
-                lossf += loss.detach()  # keep track of the mean loss
+                lossf += loss.item()  # keep track of the mean loss
             # backward pass
             if not args.inference_only:
                 loss.backward()

simple_stories_train/utils.py

@@ -4,7 +4,6 @@
 
 import torch
 import wandb
-from dotenv import load_dotenv
 from torch import nn
 
 
@@ -37,7 +36,6 @@ def save_model_and_config(save_dir: Path, model: nn.Module, step: int) -> None:
 
 
 def init_wandb(config: Any, project: str) -> None:
-    load_dotenv(override=True)
     wandb.init(
         project=project,
         config=config,