Remove lambdas in model imple to enable pickling for Zero GPU

app.py

@@ -138,7 +138,7 @@ class ParticulateApp:
                 *[None] * 9
             )
 
-    @spaces.GPU
+    @spaces.GPU(duration=10)
     def _predict_impl(
         self,
         mesh,

particulate/models.py

@@ -187,12 +187,7 @@ class Articulate3D(nn.Module):
                 nn.SiLU(),
                 nn.Linear(hidden_size * 4, 1)
             )
-            self.point_mask_decoding_func = lambda p, q: [(
-                self.point_mask_decoder(torch.cat([
-                    p.unsqueeze(2).expand(-1, -1, q.size(1), -1),  # (B, N, M, D)
-                    q.unsqueeze(1).expand(-1, p.size(1), -1, -1)   # (B, N, M, D)
-                ], dim=-1)).squeeze(-1)
-            )]
+            self.point_mask_decoding_func = self._point_mask_decoding_func_single
         else:
             self.point_mask_decoder = nn.ModuleList([
                 nn.Sequential(
@@ -202,12 +197,7 @@ class Articulate3D(nn.Module):
                 )
                 for _ in range(self.num_mask_hypotheses)
             ])
-            self.point_mask_decoding_func = lambda p, q: [(
-                self.point_mask_decoder[i](torch.cat([
-                    p.unsqueeze(2).expand(-1, -1, q.size(1), -1),  # (B, N, M, D)
-                    q.unsqueeze(1).expand(-1, p.size(1), -1, -1)   # (B, N, M, D)
-                ], dim=-1)).squeeze(-1)
-            ) for i in range(self.num_mask_hypotheses)]
+            self.point_mask_decoding_func = self._point_mask_decoding_func_multi
 
         self.use_point_features_for_motion_decoding = use_point_features_for_motion_decoding
         self.point_feature_random_ratio = point_feature_random_ratio
@@ -261,6 +251,22 @@ class Articulate3D(nn.Module):
 
         self.matcher = HungarianMatcher()
 
+    def _point_mask_decoding_func_single(self, p, q):
+        return [(
+            self.point_mask_decoder(torch.cat([
+                p.unsqueeze(2).expand(-1, -1, q.size(1), -1),  # (B, N, M, D)
+                q.unsqueeze(1).expand(-1, p.size(1), -1, -1)   # (B, N, M, D)
+            ], dim=-1)).squeeze(-1)
+        )]
+
+    def _point_mask_decoding_func_multi(self, p, q):
+        return [(
+            self.point_mask_decoder[i](torch.cat([
+                p.unsqueeze(2).expand(-1, -1, q.size(1), -1),  # (B, N, M, D)
+                q.unsqueeze(1).expand(-1, p.size(1), -1, -1)   # (B, N, M, D)
+            ], dim=-1)).squeeze(-1)
+        ) for i in range(self.num_mask_hypotheses)]
+
     def forward_attn(
         self,
         xyz: torch.FloatTensor,