correct position ids dtype

keras_hub/src/models/vit/vit_layers.py

@@ -121,8 +121,15 @@ def build(self, input_shape):
             name="position_embedding",
         )
         self.position_embedding.build([1, self.num_positions])
-        self.position_ids = ops.expand_dims(
-            ops.arange(self.num_positions), axis=0
+        self.position_ids = self.add_weight(
+            shape=(1, self.num_positions),
+            initializer="zeros",
+            # Let the backend determine the int dtype. For example, tf
+            # requires int64 for correct device placement, whereas jax and torch
+            # don't.
+            dtype=int,
+            trainable=False,
+            name="position_ids",
         )
         self.built = True