Internal change

PiperOrigin-RevId: 537247421

Author: Marvin182

clu/profiler.py

@@ -16,9 +16,9 @@
 
 Where results are stored depends on the platform (e.g. TensorBoard).
 """
-
+from collections.abc import Callable, Sequence
 import threading
-from typing import Callable, Optional, Sequence
+from typing import Any, Optional
 
 from absl import logging
 

setup.py

@@ -27,10 +27,7 @@
     "tensorflow",
     "tensorflow_datasets",
     # Temporarily disabling torch-2.0.0 since this fails Github actions.
-    "torch>=1.13.0,!=2.0.0",
-]
-pytorch_require = [
-    "torch>=1.13.0",
+    "torch>=1.13.0,<2.0.0",
 ]
 
 setup(
@@ -58,7 +55,7 @@
         "wrapt",
     ],
     tests_require=tests_require,
-    extras_require=dict(test=tests_require, pytorch=pytorch_require),
+    extras_require=dict(test=tests_require),
     classifiers=[
         "Development Status :: 4 - Beta",
         "Intended Audience :: Developers",