things

src/model/dino.py

@@ -419,6 +419,7 @@ class DINOv3ViTModel(nn.Module):
         hidden_states = self.embeddings(pixel_values, bool_masked_pos=bool_masked_pos)
         position_embeddings = self.rope_embeddings(pixel_values)
 
+        latents = []
         for i, layer_module in enumerate(self.layers):
             layer_head_mask = head_mask[i] if head_mask is not None else None
             hidden_states = layer_module(
@@ -426,11 +427,13 @@ class DINOv3ViTModel(nn.Module):
                 attention_mask=layer_head_mask,
                 position_embeddings=position_embeddings,
             )
+            latents.append(hidden_states)
 
         sequence_output = self.norm(hidden_states)
         pooled_output = sequence_output[:, 0, :]
 
         return {
             "last_hidden_state": sequence_output,
+            "latents": latents,
             "pooler_output": pooled_output,
         }