Change from return type[] to return Type[]

Author: tvmarino

compiler_opt/rl/inlining/imitation_learning_config.py

@@ -14,6 +14,7 @@
 """Module for collect data of inlining-for-size."""
 
 import gin
+from typing import Type
 
 import numpy as np
 import tensorflow as tf
@@ -75,7 +76,7 @@ def get_input_signature():
 
 
 @gin.register
-def get_task_type() -> type[env.InliningForSizeTask]:
+def get_task_type() -> Type[env.InliningForSizeTask]:
   """Returns the task type for the trajectory collection."""
   return env.InliningForSizeTask