updata main.py

main.py

@@ -141,7 +141,7 @@ def train_model(
             #print(pruner.group_reg[pruner._groups[0]])
         
         model.eval()
-        acc, val_loss = eval(model, test_loader, device=device)
+        acc, val_loss, eval_time = eval(model, test_loader, device=device)
         args.logger.info(
             "Epoch {:d}/{:d}, Acc={:.4f}, Val Loss={:.4f}, lr={:.4f}".format(
                 epoch, epochs, acc, val_loss, optimizer.param_groups[0]["lr"]
@@ -297,13 +297,13 @@ def main():
         # 1. Pruning
         model.eval()
         ori_ops, ori_size = tp.utils.count_ops_and_params(model, example_inputs=example_inputs)
-        ori_acc, ori_val_loss = eval(model, test_loader, device=args.device)
+        ori_acc, ori_val_loss, ori_eval_time = eval(model, test_loader, device=args.device)
         args.logger.info("Pruning...")
         progressive_pruning(pruner, model, speed_up=args.speed_up, example_inputs=example_inputs)
         del pruner # remove reference
         args.logger.info(model)
         pruned_ops, pruned_size = tp.utils.count_ops_and_params(model, example_inputs=example_inputs)
-        pruned_acc, pruned_val_loss = eval(model, test_loader, device=args.device)
+        pruned_acc, pruned_val_loss, pruned_eval_time = eval(model, test_loader, device=args.device)
         
         args.logger.info(
             "Params: {:.2f} M => {:.2f} M ({:.2f}%)".format(