From 507ba6ae90245bc23a88d0cf8301b262da4ec7bc Mon Sep 17 00:00:00 2001
From: Neel Gupta <neelgupta04@outlook.com>
Date: Tue, 4 Jun 2024 12:44:02 +0100
Subject: [PATCH] chore: Update HF token and LR hyperparam schedule

---
 Dockerfile     | 2 --
 train_model.py | 2 +-
 2 files changed, 1 insertion(+), 3 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index bb8e0c6..9f86538 100755
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,8 +1,6 @@
 FROM python:3.11
 
 # Set environment variables
-ENV WANDB_API_KEY=78c7285b02548bf0c06dca38776c08bb6018593f
-ENV HF_TOKEN=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
 ENV jax_threefry_partitionable=1
 
 # Install base utilities
diff --git a/train_model.py b/train_model.py
index 83cebf1..c1b9d32 100644
--- a/train_model.py
+++ b/train_model.py
@@ -136,7 +136,7 @@ def kickoff_optuna(trial, **trainer_kwargs):
     args.epochs = 1
     
     # Regularization hyperparams
-    args.lr = trial.suggest_float('lr', 1e-4, 3e-2)
+    args.lr = trial.suggest_float('lr', 1e-4, 9e-2)
     args.drop_rate = trial.suggest_float('drop_rate', 0.0, 0.1, step=0.01)
     args.weight_decay = trial.suggest_float('weight_decay', 1e-5, 1e-3)
     args.warmup_steps = trial.suggest_int('warmup_steps', 0, 500, step=100)