diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index 2f7c5ef5..d6461f3b 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -46,7 +46,9 @@ jobs:
         task: [Typecheck, Lint, Ruff, Yapf, Test]
         include:
           - task: Typecheck
-            cmd: pytype -j auto .
+            # We disable pyi-error due to https://github.com/google/pytype/issues/764
+            # and other related issues with tensorflow imports.
+            cmd: pytype -j auto --disable=pyi-error .
           - task: Lint
             cmd: pylint --rcfile .pylintrc --recursive yes .
           - task: Ruff
diff --git a/compiler_opt/es/blackbox_learner.py b/compiler_opt/es/blackbox_learner.py
index 597affd1..80dc9495 100644
--- a/compiler_opt/es/blackbox_learner.py
+++ b/compiler_opt/es/blackbox_learner.py
@@ -31,10 +31,6 @@
 from compiler_opt.rl import policy_saver
 from compiler_opt.es import blackbox_evaluator  # pylint: disable=unused-import
 
-# Pytype cannot pick up the pyi file for tensorflow.summary. Disable the error
-# here as these errors are false positives.
-# pytype: disable=pyi-error
-
 # If less than 40% of requests succeed, skip the step.
 _SKIP_STEP_SUCCESS_RATIO = 0.4
 
diff --git a/compiler_opt/rl/compilation_runner_test.py b/compiler_opt/rl/compilation_runner_test.py
index ad9b0461..abef9fb2 100644
--- a/compiler_opt/rl/compilation_runner_test.py
+++ b/compiler_opt/rl/compilation_runner_test.py
@@ -23,8 +23,7 @@
 from absl import flags
 import tensorflow as tf
 
-# This is https://github.com/google/pytype/issues/764
-from google.protobuf import text_format  # pytype: disable=pyi-error
+from google.protobuf import text_format
 
 from compiler_opt.rl import compilation_runner
 from compiler_opt.rl import constant
diff --git a/compiler_opt/rl/imitation_learning/generate_bc_trajectories_test.py b/compiler_opt/rl/imitation_learning/generate_bc_trajectories_test.py
index 1414eb90..4ec0de67 100644
--- a/compiler_opt/rl/imitation_learning/generate_bc_trajectories_test.py
+++ b/compiler_opt/rl/imitation_learning/generate_bc_trajectories_test.py
@@ -27,7 +27,7 @@
 from tf_agents.trajectories import time_step
 from tf_agents.system import system_multiprocessing as multiprocessing
 
-from google.protobuf import text_format  # pytype: disable=pyi-error
+from google.protobuf import text_format
 
 from compiler_opt.rl.imitation_learning import generate_bc_trajectories_lib
 from compiler_opt.rl import env
diff --git a/compiler_opt/rl/imitation_learning/weighted_bc_trainer_lib.py b/compiler_opt/rl/imitation_learning/weighted_bc_trainer_lib.py
index e1379bd1..221d3513 100644
--- a/compiler_opt/rl/imitation_learning/weighted_bc_trainer_lib.py
+++ b/compiler_opt/rl/imitation_learning/weighted_bc_trainer_lib.py
@@ -43,10 +43,6 @@
     ('Directory containing the quantile map for normalizing features'
      'in feature_ops.build_quantile_map.'))
 
-# Pytype cannot pick up the pyi file for tensorflow.summary. Disable the error
-# here as these errors are false positives.
-# pytype: disable=pyi-error
-
 
 @gin.configurable
 class TrainingWeights:
diff --git a/compiler_opt/rl/local_data_collector_test.py b/compiler_opt/rl/local_data_collector_test.py
index 95c42b7b..4f866cca 100644
--- a/compiler_opt/rl/local_data_collector_test.py
+++ b/compiler_opt/rl/local_data_collector_test.py
@@ -20,8 +20,7 @@
 import tensorflow as tf
 from tf_agents.system import system_multiprocessing as multiprocessing
 
-# This is https://github.com/google/pytype/issues/764
-from google.protobuf import text_format  # pytype: disable=pyi-error
+from google.protobuf import text_format
 from compiler_opt.distributed.local.local_worker_manager import LocalWorkerPoolManager
 from compiler_opt.rl import compilation_runner
 from compiler_opt.rl import corpus
diff --git a/compiler_opt/rl/log_reader_test.py b/compiler_opt/rl/log_reader_test.py
index db5791ea..b5dcc48b 100644
--- a/compiler_opt/rl/log_reader_test.py
+++ b/compiler_opt/rl/log_reader_test.py
@@ -18,8 +18,7 @@
 import json
 from compiler_opt.rl import log_reader
 
-# This is https://github.com/google/pytype/issues/764
-from google.protobuf import text_format  # pytype: disable=pyi-error
+from google.protobuf import text_format
 from typing import BinaryIO
 
 import numpy as np
diff --git a/compiler_opt/rl/train_bc.py b/compiler_opt/rl/train_bc.py
index 7f763a98..695d9ad0 100644
--- a/compiler_opt/rl/train_bc.py
+++ b/compiler_opt/rl/train_bc.py
@@ -20,10 +20,6 @@
 from absl import logging
 import gin
 
-# Pytype cannot pick up the pyi file for tensorflow.summary. Disable the error
-# here as these errors are false positives.
-# pytype: disable=pyi-error
-
 # <Internal> Using XM - flags.  # pylint: disable=unused-import
 from compiler_opt.rl import agent_config
 from compiler_opt.rl import data_reader
diff --git a/compiler_opt/tools/generate_default_trace_test.py b/compiler_opt/tools/generate_default_trace_test.py
index 0ea03b32..4abafa75 100644
--- a/compiler_opt/tools/generate_default_trace_test.py
+++ b/compiler_opt/tools/generate_default_trace_test.py
@@ -22,8 +22,7 @@
 import gin
 import tensorflow as tf
 
-# This is https://github.com/google/pytype/issues/764
-from google.protobuf import text_format  # pytype: disable=pyi-error
+from google.protobuf import text_format
 from compiler_opt.rl import compilation_runner
 from compiler_opt.tools import generate_default_trace