Update app.py

app.py

@@ -29,14 +29,16 @@ class TransformNet(nn.Module):
 
     def forward(self, x):
         if x.dim() == 5:  
-            x = x.squeeze(-1)  
-        
+            x = x.squeeze(-1)  # ✅ Remove extra dimension if exists
+    
         x = self.conv1(x)
         x = self.conv2(x)
-        x = self.conv3(x)
+    
+        x = x.squeeze(-1)  # ✅ Ensure shape is [batch, channels, length] before Conv1d
+        x = self.conv3(x)  # ✅ Now Conv1d receives correct input shape [batch, channels, length]
+    
         return self.fc(x.max(dim=-1)[0])  # ✅ Ensure correct pooling
 
-
 class PointCloudEncoder(nn.Module):
     """Point Cloud Encoder (pc_enc)"""
     def __init__(self):
@@ -61,16 +63,19 @@ class PointCloudEncoder(nn.Module):
         ])
         self.lin_global = nn.Linear(64, 128)
 
-    def forward(self, x):
+   def forward(self, x):
         x = self.transform_net(x)
+    
         for i, conv in enumerate(self.convs):
             if i < 2:
-                x = conv(x)  # ✅ Conv2d
+                x = conv(x)  # ✅ Conv2d keeps 4D
             else:
-                x = x.squeeze(-1)  # ✅ Remove extra dimension before Conv1d
-                x = conv(x)  # ✅ Conv1d (Matches `[64, 128, 1]`)
+                x = x.squeeze(-1)  # ✅ Ensure shape is [batch, channels, length] before Conv1d
+                x = conv(x)  # ✅ Conv1d now works with the correct input
+    
         return self.lin_global(x.max(dim=-1)[0])  # ✅ Fix pooling
 
+
 class GarmentEncoder(nn.Module):
     """Garment Feature Encoder (garm_enc)"""
     def __init__(self, num_classes=18, feature_dim=64):