Update app.py

app.py

@@ -354,7 +354,14 @@ class ForgeryDetector:
         # Classify regions
         results = []
         for region in regions:
-            df = decoder_features[0]          # shape: [C, H, W]
+            # Get decoder features and handle shape
+            df = decoder_features[0].cpu()  # Get first decoder feature
+            
+            # Remove batch dimension if present: [1, C, H, W] -> [C, H, W]
+            if df.ndim == 4:
+                df = df.squeeze(0)
+            
+            # Now df should be [C, H, W]
             _, fh, fw = df.shape
 
             region_mask = region['region_mask']