Update trainer.py

trainer.py

@@ -16,9 +16,6 @@ from tqdm import tqdm
 from torchvision import datasets, transforms
 from torch.utils.tensorboard import SummaryWriter
 
-from tri_stream_vit import (
-    create_tri_stream_vit, MasteryQueue, SimplexBuffer)
-
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 torch.backends.cuda.matmul.allow_tf32 = True
 torch.backends.cudnn.allow_tf32 = True
@@ -170,10 +167,10 @@ print(f"  GAL Procrustes: every {GAL_UPDATE_INTERVAL} batches, "
       f"lr={GAL_LR}, whiten={USE_WHITENED_PROCRUSTES}")
 print(f"  Mastery: patience={MASTERY_PATIENCE}, queue adaptive "
       f"[{MASTERY_MIN_SIZE}–{MASTERY_MAX_SIZE}]")
-print(f"  Optimizer: AdamW (wd=0.01)")
+print(f"  Optimizer: Adam")
 print(f"{'='*60}")
 
-optimizer = torch.optim.AdamW(model.parameters(), lr=LR, weight_decay=0.01)
+optimizer = torch.optim.Adam(model.parameters(), lr=LR)
 
 total_steps = len(train_loader) * EPOCHS
 warmup_steps = len(train_loader) * WARMUP