diff --git a/examples/mnist/mnist.zig b/examples/mnist/mnist.zig
index 632cafd..4130bf7 100644
--- a/examples/mnist/mnist.zig
+++ b/examples/mnist/mnist.zig
@@ -27,7 +27,7 @@ const Mnist = struct {
     /// just two linear layers + relu activation
     pub fn forward(self: Mnist, input: zml.Tensor) zml.Tensor {
         // std.log.info("Compiling for target: {s}", .{@tagName(input.getContext().target())});
-        var x = input.flattenAll().convert(.f32);
+        var x = input.flatten().convert(.f32);
         const layers: []const Layer = &.{ self.fc1, self.fc2 };
         for (layers) |layer| {
             x = layer.forward(x);
diff --git a/examples/modernbert/modernbert.zig b/examples/modernbert/modernbert.zig
index 2e6eb4a..1ba3557 100644
--- a/examples/modernbert/modernbert.zig
+++ b/examples/modernbert/modernbert.zig
@@ -43,7 +43,7 @@ pub const ModernBertForMaskedLM = struct {
         const biased_logits = logits.add(self.decoder.bias.withTags(.{.voc}).broad(logits.shape()));
 
         const probabilities = biased_logits.softmax(.voc);
-        return probabilities.topK(5, .voc, .{ .descending = true });
+        return probabilities.topK(.{ .best_words = .voc }, 5, .{ .descending = true });
     }
 };