Fixes unit test timeouts.

PiperOrigin-RevId: 424332493
Change-Id: I54c0de7b8fd5f701f9743552f5b594d91acc593f

Author: Johannes Ballé

tensorflow_compression/python/entropy_models/BUILD

@@ -51,9 +51,10 @@ py_library(
 
 py_test(
     name = "continuous_indexed_test",
+    timeout = "long",
     srcs = ["continuous_indexed_test.py"],
     python_version = "PY3",
-    shard_count = 5,
+    shard_count = 6,
     deps = [
         ":continuous_indexed",
         "//tensorflow_compression/python/distributions:uniform_noise",