diff --git a/learning_to_simulate/graph_network.py b/learning_to_simulate/graph_network.py
index 92132fa..de38c59 100644
--- a/learning_to_simulate/graph_network.py
+++ b/learning_to_simulate/graph_network.py
@@ -45,7 +45,7 @@ def build_mlp(
     hidden_size: int, num_hidden_layers: int, output_size: int) -> snt.Module:
   """Builds an MLP."""
   return snt.nets.MLP(
-      output_sizes=[hidden_size] * num_hidden_layers + [output_size])
+      output_sizes=[hidden_size] * num_hidden_layers + [output_size])  # pytype: disable=bad-return-type  # gen-stub-imports
 
 
 class EncodeProcessDecode(snt.AbstractModule):
diff --git a/mmv/utils/ucf101_dataset.py b/mmv/utils/ucf101_dataset.py
index 41304b6..63c5af8 100644
--- a/mmv/utils/ucf101_dataset.py
+++ b/mmv/utils/ucf101_dataset.py
@@ -58,7 +58,7 @@ class ModUcf101(tfds.video.Ucf101):
     features = tfds.features.FeaturesDict({
         'video': tfds.features.Video(video_shape,
                                      ffmpeg_extra_args=ffmpeg_extra_args,
-                                     encoding_format='jpeg'),
+                                     encoding_format='jpeg'),  # pytype: disable=wrong-arg-types  # gen-stub-imports
         'label': tfds.features.ClassLabel(names_file=labels_names_file),
     })
     return tfds.core.DatasetInfo(