fix pixel_values being casted to int

vit_gpt2/modeling_flax_vit_gpt2_lm.py

@@ -241,7 +241,7 @@ class FlaxViTGPT2LMPreTrainedModel(FlaxPreTrainedModel):
 
         return self.module.apply(
             {"params": params or self.params},
-            pixel_values=jnp.array(pixel_values, dtype="i4"),
+            pixel_values=jnp.array(pixel_values, dtype=jnp.float32),
             output_attentions=output_attentions,
             output_hidden_states=output_hidden_states,
             return_dict=return_dict,