diff --git a/eval/UnlearningAttack.py b/eval/UnlearningAttack.py
index 7de536d..87bbb02 100644
--- a/eval/UnlearningAttack.py
+++ b/eval/UnlearningAttack.py
@@ -101,7 +101,7 @@ class UnlearningAttack:
             self.hook = None
 
     def _extract_attack_features(self, target_model, loader, device, target_class):
-        target_model.eval()
+        '''target_model.eval()
         all_probs = []
         all_entropies = []
         all_losses = []
@@ -142,10 +142,39 @@ class UnlearningAttack:
         else:
             compiled_latent = np.zeros((len(X_features), 512))
 
-        return X_features, compiled_latent
+        return X_features, compiled_latent'''
+        target_model.eval()
+        all_probs, all_entropies, all_losses = [], [], []
+        
+        with torch.no_grad():
+            for data, targets in loader:
+                data, targets = data.to(device), targets.to(device)
+                if target_model.__class__.__name__ == "WF_Module":
+                    gate = torch.full((data.size(0),), target_class, dtype=torch.long, device=device)
+                    outputs = target_model(data, target_class_indices=gate)
+                else:
+                    outputs = target_model(data)
+                
+                probs = torch.softmax(outputs, dim=1)
+                all_probs.extend(probs.cpu().numpy())
+                
+                # Entropy: -sum(p * log(p))
+                log_probs = torch.log(probs + 1e-10)
+                entropy = -torch.sum(probs * log_probs, dim=1)
+                all_entropies.extend(entropy.cpu().numpy())
+                
+                loss = self.criterion(outputs, targets)
+                all_losses.extend(loss.cpu().numpy())
+
+        # Combine output-based features only
+        return np.hstack([
+            np.array(all_probs), 
+            np.array(all_entropies).reshape(-1, 1), 
+            np.array(all_losses).reshape(-1, 1)
+        ])
 
     def run_parameter_space_mia(self, unlearned_model, shadow_model, forget_loader, retain_test_loader, device, index):
-        X_shadow_mem, _ = self._extract_attack_features(shadow_model, forget_loader, device, index)
+        '''X_shadow_mem, _ = self._extract_attack_features(shadow_model, forget_loader, device, index)
         X_shadow_non, _ = self._extract_attack_features(shadow_model, retain_test_loader, device, index)
         
         min_train = min(len(X_shadow_mem), len(X_shadow_non))
@@ -165,10 +194,34 @@ class UnlearningAttack:
         predictions = attack_classifier.predict(X_test)
         mia_accuracy = accuracy_score(y_test, predictions)
         
-        clean_centroid = np.mean(retain_latent, axis=0)
-        forget_distances = np.linalg.norm(latent_features - clean_centroid, axis=1)
+        #clean_centroid = np.mean(retain_latent, axis=0)
+        #forget_distances = np.linalg.norm(latent_features - clean_centroid, axis=1)
         
-        return mia_accuracy, np.mean(forget_distances)
+        return mia_accuracy, np.mean(forget_distances)'''
+        X_shadow_mem = self._extract_attack_features(shadow_model, forget_loader, device, index)
+        X_shadow_non = self._extract_attack_features(shadow_model, retain_test_loader, device, index)
+        
+        # Train MIA Classifier
+        min_train = min(len(X_shadow_mem), len(X_shadow_non))
+        X_train = np.vstack([X_shadow_mem[:min_train], X_shadow_non[:min_train]])
+        y_train = np.concatenate([np.ones(min_train), np.zeros(min_train)])
+        
+        attack_classifier = LogisticRegression(max_iter=1000)
+        attack_classifier.fit(X_train, y_train)
+        
+        # Evaluate MIA
+        X_eval_mem = self._extract_attack_features(unlearned_model, forget_loader, device, index)
+        X_eval_non = self._extract_attack_features(unlearned_model, retain_test_loader, device, index)
+        
+        min_test = min(len(X_eval_mem), len(X_eval_non))
+        X_test = np.vstack([X_eval_mem[:min_test], X_eval_non[:min_test]])
+        y_test = np.concatenate([np.ones(min_test), np.zeros(min_test)])
+        
+        predictions = attack_classifier.predict(X_test)
+        mia_accuracy = accuracy_score(y_test, predictions)
+        
+        # Note: Latent distance is removed as it's not a black-box metric
+        return mia_accuracy, 0.0
     
     def _comput_adversarial_accuracy(self, filtered, naive, axis=-1):
 
@@ -282,11 +335,16 @@ class UnlearningAttack:
 
         # 2. Extract latent features for A-Dist
         # We need features from both Unlearned and Retrained model
-        _, unlearned_latent = self._extract_attack_features(unlearned_instance.model, forget_loader, device, target_class)
-        _, retrained_latent = self._extract_attack_features(reference_model_torch, forget_loader, device, target_class)
-        
+        #_, unlearned_latent = self._extract_attack_features(unlearned_instance.model, forget_loader, device, target_class)
+        #_, retrained_latent = self._extract_attack_features(reference_model_torch, forget_loader, device, target_class)
+        # Extract features (now just one returned object)
+        unlearned_features = self._extract_attack_features(unlearned_instance.model, forget_loader, device, target_class)
+        retrained_features = self._extract_attack_features(reference_model_torch, forget_loader, device, target_class)
+
+        # Calculate A-Dist using these features
+        a_dist = self.calculate_a_dist(unlearned_features, retrained_features)
         # 3. Calculate A-Dist (Replacing latent_distance)
-        a_dist = self.calculate_a_dist(unlearned_latent, retrained_latent)
+        #a_dist = self.calculate_a_dist(unlearned_latent, retrained_latent)
 
 
         print(f"[{framework_name}] Class {target_class} | Parameter MIA: {parameter_mia_acc:.4f} | Latent Dist: {latent_dist:.4f} | Lookalike: {lookalike_acc:.4f}" )