Testing

test.py

@@ -34,7 +34,7 @@ class SkinGPTTester:
                 return
             
             # Get the last layer's attention
-            attention = self.classifier.model.q_former.last_attention[0].mean(dim=0)
+            attention = self.classifier.model.q_former.last_attention[0][0]  # shape: [num_tokens,]
             
             # Print attention shape for debugging
             print(f"Attention shape: {attention.shape}")
@@ -124,7 +124,7 @@ class SkinGPTTester:
         # Extract visual features using attention
         with torch.no_grad():
             image_tensor = self.transform(image).unsqueeze(0).to(self.device)
-            attention = self.classifier.model.q_former.last_attention[0].mean(dim=0)
+            attention = self.classifier.model.q_former.last_attention[0][0]
         
         # Get regions with high attention
         attention = attention.reshape(int(math.sqrt(attention.shape[1])), -1)