Updating prepare_models (#1)

- Updating prepare_models (8e98ce788746309169c34d162fa96fa2d4d7de56)

util/prepare_utils.py

@@ -184,17 +184,14 @@ def extract_features(imgs, feature_extractor_ensemble, dim):
     return features
 
 
-
-
-
 def prepare_models(model_backbones,
-                 input_size,
-                 model_roots,
-                 kernel_size_attack,
-                 sigma_attack,
-                 combination,
-                 using_subspace,
-                 V_reduction_root):
+                   input_size,
+                   model_roots,
+                   kernel_size_attack,
+                   sigma_attack,
+                   combination,
+                   using_subspace,
+                   V_reduction_root):
 
     backbone_dict = {'IR_50': IR_50(input_size), 'IR_152': IR_152(input_size), 'ResNet_50': ResNet_50(input_size),
                      'ResNet_152': ResNet_152(input_size)}
@@ -205,11 +202,14 @@ def prepare_models(model_backbones,
     models_attack = []
     for i in range(len(model_backbones)):
         model = backbone_dict[model_backbones[i]]
-        model.load_state_dict(torch.load(model_roots[i], map_location=device))
+        try:
+            model.load_state_dict(torch.load(model_roots[i], map_location=device))
+        except Exception as e:
+            print(f"Error loading model {model_roots[i]}: {e}")
+            continue
         models_attack.append(model)
 
     if using_subspace:
-
         V_reduction = []
         for i in range(len(model_backbones)):
             V_reduction.append(torch.tensor(np.load(V_reduction_root[i])))
@@ -221,6 +221,7 @@ def prepare_models(model_backbones,
 
     return models_attack, V_reduction, dim
 
+
 def prepare_data(query_data_root, target_data_root, freq, batch_size, warp = False, theta_warp = None):
 
     data = datasets.ImageFolder(query_data_root, tensor_transform)