diff --git a/examples/mnist/mnist.zig b/examples/mnist/mnist.zig
index 0f35c7e..51ca1e4 100644
--- a/examples/mnist/mnist.zig
+++ b/examples/mnist/mnist.zig
@@ -33,7 +33,7 @@ const Mnist = struct {
         for (layers) |layer| {
             x = zml.call(layer, .forward, .{x});
         }
-        return x.argMax(0, .u8).indices;
+        return x.argMax(0).indices.convert(.u8);
     }
 };