Upload train_shakespeare.py

train_shakespeare.py

@@ -6,14 +6,11 @@ from dataclasses import dataclass
 import torch
 import torch.nn as nn
 from torch.nn import functional as F
-import wandb
 
 # Set MPS memory management
 os.environ['PYTORCH_MPS_HIGH_WATERMARK_RATIO'] = '0.0'
 os.environ['PYTORCH_MPS_LOW_WATERMARK_RATIO'] = '0.5'
 
-# Initialize wandb
-wandb.init(project="shakespeare-gpt", name="gpt2-124M-training")
 
 class CausalSelfAttention(nn.Module):
     def __init__(self, config):
@@ -204,10 +201,6 @@ for iter in range(num_iters):
     if iter % eval_interval == 0:
         current_loss = loss.item()
         print(f'step {iter}, loss: {current_loss:.4f}')
-        wandb.log({
-            "iter": iter,
-            "loss": current_loss
-        })
         
         # Save if this is the best model so far
         if current_loss < best_loss:
@@ -243,5 +236,3 @@ torch.save({
     'loss': loss.item(),
     'best_loss': best_loss,
 }, final_path)
-
-wandb.finish()