Fix top_k usage

keras_hub/src/models/deepseek_r1/deepseek_backbone.py

@@ -71,7 +71,7 @@ class ModelArgs:
     inter_dim: int = 10944
     moe_inter_dim: int = 1408
     # n_layers: int = 27
-    n_layers: int = 1
+    n_layers: int = 2
     n_dense_layers: int = 1
     n_heads: int = 16
     # moe

keras_hub/src/models/deepseek_r1/deepseek_layers.py

@@ -577,7 +577,7 @@ def call(self, x):
 
             scores = ops.where(ops.expand_dims(mask, -1), scores, float("inf"))
             scores = ops.reshape(scores, [scores.shape[0], -1])
-        indices = ops.top_k(scores, self.topk, dim=-1)[1]
+        indices = ops.top_k(scores, self.topk)[1]
         weights = ops.take_along_axis(original_scores, indices, axis=1)
         if self.score_func == "sigmoid":
             weights /= ops.sum(weights, axis=-1, keepdims=True)

tools/checkpoint_conversion/convert_deepseek_checkpoints.py

@@ -218,7 +218,7 @@ def convert_weights():
 
     n_blocks = len(model.layers) - 3  # (3 = len(embed, head, norm))
     for i in range(n_blocks):
-        convert_block(model.layers[i + 1], torch_weights, i + 1)
+        convert_block(model.layers[i + 2], torch_weights, i + 1)
 
     # Run some tokens as a sanity check
     total_tokens_generated = 0