Update scrub.py

scrub.py

@@ -344,12 +344,13 @@ def unlearn():
     model.load_state_dict(torch.load('faces_best_model.pth', map_location=torch.device('cpu')))
     model_eval = copy.deepcopy(model)
     model_eval.eval()
+    print("BEGIN INTIALZINGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGG")
     forget_class = get_forget_class('forget_set', model_eval)
     mean, std, im_size = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225], 224
     tfs = transforms.Compose([transforms.Resize((im_size, im_size)), transforms.ToTensor(), transforms.Normalize(mean = mean, std = std)])
     will_tr_dl, will_val_dl, will_ts_dl, classes = get_dls(root = "forget_set", forget_class=forget_class, transformations = tfs, bs = 32, single=True)
     celebs_tr_dl, celebs_val_dl, celebs_ts_dl, classes = get_dls(root = "celeb-dataset", forget_class=forget_class, transformations = tfs, bs = 32)
-
+    print("BEGIN PEeEEEEEEEEEEEEEEEEEEEEEEERPARING FOR UNLEARNINGGGGGGGG")
     args = Args()
     args.optim = 'sgd'
     args.gamma = 0.99
@@ -395,8 +396,9 @@ def unlearn():
                               weight_decay=args.sgda_weight_decay)
     
     module_list.append(model_t)
-    
+    print("BEGIN UNLEARNINGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGG")
     for epoch in tqdm(range(1, args.sgda_epochs + 1)):
+        print("\n\n==============================>epoch: ", epoch)
         maximize_loss = 0
         if epoch <= args.msteps:
             maximize_loss = train_distill(epoch, will_tr_dl, module_list, swa_model, criterion_list, optimizer, args, "maximize")