Upload model

modeling.py

@@ -433,7 +433,7 @@ class LSPTransformer(nn.Module):
     def init_weights(self) -> None:
         prior_prob = 0.01
         bias_value = -math.log((1 - prior_prob) / prior_prob)
-        self.class_head.bias.data = torch.ones(self.num_classes) * bias_value
+        nn.init.constant_(self.class_head.bias, bias_value)
 
         # initialize regression layers
         for head in self.point_head:
@@ -519,6 +519,19 @@ class LSPTransformer(nn.Module):
                 ref_points, self.query_block_size, self.query_block_size
             ).flatten(1, 2),
             "embeddings": tgt.flatten(1, 2),
+            "aux_outputs": [
+                {
+                    "logits": a,
+                    "points": b,
+                    "radial_distances": c,
+                }
+                for a, b, c in zip(
+                    logits_list[:-1],
+                    ref_points_list[:-1],
+                    radial_distances_list[:-1],
+                    strict=True,
+                )
+            ],
         }