diff --git a/compiler_opt/rl/env.py b/compiler_opt/rl/env.py
index 1bd96afb..904fd388 100644
--- a/compiler_opt/rl/env.py
+++ b/compiler_opt/rl/env.py
@@ -359,6 +359,7 @@ def reset(self, module: corpus.LoadedModuleSpec):
     self._clang_generator.send(None)
     # pytype: disable=attribute-error
     self._iclang, self._clang = self._clang_generator.send(module)
+    # pytype: enable=attribute-error
     return self._get_observation()
 
   def step(self, action: np.ndarray):