fix type

Author: wenyi-guo

keras_rs/src/utils/tpu_test_utils.py

@@ -34,7 +34,7 @@ def experimental_distribute_dataset(
 
 class JaxDummyStrategy(DummyStrategy):
     @property
-    def num_replicas_in_sync(self) -> int:
+    def num_replicas_in_sync(self) -> Any:
         if jax is None:
             return 0
         return jax.device_count("tpu")