diff --git a/tensorflow_graphics/projects/gan/architectures_style_gan_v2.py b/tensorflow_graphics/projects/gan/architectures_style_gan_v2.py
index ab0763bcf..2c917d978 100644
--- a/tensorflow_graphics/projects/gan/architectures_style_gan_v2.py
+++ b/tensorflow_graphics/projects/gan/architectures_style_gan_v2.py
@@ -85,7 +85,7 @@ def get_noise_dimensions(
   return noise_dimensions
 
 
-def _create_noise_inputs(num_upsampling_blocks: int) -> List[tf.keras.Input]:  # pytype: disable=invalid-annotation  # typed-keras
+def _create_noise_inputs(num_upsampling_blocks: int) -> List[tf.Tensor]:
   """Creates the noise input layers."""
   noise_dimensions = get_noise_dimensions(num_upsampling_blocks)
   noise_inputs = [