Update data split to 40%

Author: neel04

train_model.py

@@ -50,9 +50,9 @@ def main(key: PRNGKeyArray):
 
     if args.tune_hyperparams:
         args.group = 'Sweeps' if args.baseline else 'Sweeps_5i'
-        
-        trainloader = train_dataset.create_dataloader('20%')
-        valloader = val_dataset.create_dataloader('20%')
+
+        trainloader = train_dataset.create_dataloader("40%")
+        valloader = val_dataset.create_dataloader("40%")
 
         trainloader = train_dataset.create_dataloader("40%")
         valloader = val_dataset.create_dataloader("40%")
@@ -136,7 +136,7 @@ def main(key: PRNGKeyArray):
 def kickoff_optuna(trial, **trainer_kwargs):
     args = trainer_kwargs['args']
 
-    args.epochs = 2
+    args.epochs = 1
 
     args.lr = trial.suggest_float('lr', 1e-4, 1e-2, step=1e-4)
     args.drop_rate = trial.suggest_float('drop_rate', 0.0, 0.1, step=0.01)