next

stargan/attacks.py

@@ -24,7 +24,7 @@ class LinfPGDAttack(object):
         X = X_nat.clone().detach_()
 
         for i in range(self.k):
-            print(i)
+            # print(i)
             X.requires_grad = True
             output, feats = self.model(X, c_trg)
 

stargan/solver.py

@@ -594,6 +594,7 @@ class Solver(object):
         n_samples = 0
 
         for i, (x_real, c_org) in enumerate(data_loader):
+            self.restore_model(self.test_iters)
             # Black image
             black = np.zeros((1,3,256,256))
             black = torch.FloatTensor(black).to(self.device)