fixed error

main.py

@@ -109,7 +109,7 @@ def test(model, dataloader):
     return correct
 
 
-def forward_epoch(model, optimizer, state_keeper, time, epochs):
+def forward_epoch(model, train_dataloader, test_dataloader, optimizer, state_keeper, time, epochs):
     for epoch in range(1, epochs + 1):
         loss_dict = train(model, optimizer, train_dataloader)
         with torch.no_grad():
@@ -131,11 +131,12 @@ def forward_epoch(model, optimizer, state_keeper, time, epochs):
     for time in range(args.times):
         model = utils.get_model(args)
         optimizer = utils.get_optimizer(args.optim, args.lr, model)
-        forward_epoch(model, optimizer, state_keeper, time, args.epochs)
+        forward_epoch(model, train_dataloader, test_dataloader,
+                      optimizer, state_keeper, time, args.epochs)
         if args.exname == "TransferLearning":
             optimizer_aux = utils.get_optimizer(
                 args.optim, args.lr_aux, model)
-            forward_epoch(model, optimizer_aux, state_keeper_aux,
+            forward_epoch(model, train_dataloader_aux, test_dataloader_aux, optimizer_aux, state_keeper_aux,
                           time, args.epochs_aux)
 
     state_keeper.save()

train.sh

@@ -9,5 +9,9 @@ python main.py --batch_size 128 --lr 1e-4 --epochs 20 --times 5 --data_root data
 
 # Transfer Learning
 # MNIST -> SVHN
-python main.py --batch_size 128 --lr 1e-2 --lr_aux 1e-5 --epochs 20 --epochs_aux 200 --times 5 --data_root data --dataset MNIST --dataset_aux SVHN --num_workers 2 --net ConvMNIST --af all --optim SGD --exname TransferLearning
-python main.py --batch_size 128 --lr 1e-2 --lr_aux 1e-5 --epochs 20 --epochs_aux 200 --times 5 --data_root data --dataset SVHN --dataset_aux MNIST --num_workers 2 --net ConvMNIST --af all --optim SGD --exname TransferLearning
+python main.py --batch_size 128 --lr 1e-2 --lr_aux 1e-5 --epochs 5 --epochs_aux 100 --times 5 --data_root data --dataset MNIST --dataset_aux SVHN --num_workers 2 --net ConvMNIST --af all --optim SGD --exname TransferLearning
+python main.py --batch_size 128 --lr 1e-2 --lr_aux 1e-5 --epochs 5 --epochs_aux 100 --times 5 --data_root data --dataset SVHN --dataset_aux MNIST --num_workers 2 --net ConvMNIST --af all --optim SGD --exname TransferLearning
+
+
+python main.py --batch_size 128 --lr 1e-2 --lr_aux 1e-5 --epochs 10 --epochs_aux 100 --times 5 --data_root data --dataset MNIST --dataset_aux SVHN --num_workers 2 --net ConvMNIST --af all --optim SGD --exname TransferLearning
+python main.py --batch_size 128 --lr 1e-2 --lr_aux 1e-5 --epochs 10 --epochs_aux 100 --times 5 --data_root data --dataset SVHN --dataset_aux MNIST --num_workers 2 --net ConvMNIST --af all --optim SGD --exname TransferLearning