Fix missing group norm in converted TF Lite models. (#102)

* Fix missing group norm in converted TF Lite models. * update
google-ai-edge · Jul 24, 2024 · b178852 · b178852
1 parent 53fa236
commit b178852
Showing 1 changed file with 6 additions and 9 deletions.
diff --git a/ai_edge_torch/generative/layers/unet/blocks_2d.py b/ai_edge_torch/generative/layers/unet/blocks_2d.py
@@ -125,10 +125,9 @@ def forward(self, input_tensor: torch.Tensor) -> torch.Tensor:
     """
     residual = input_tensor
     B, C, H, W = input_tensor.shape
-    x = input_tensor
     if self.config.normalization_config.type == layers_cfg.NormalizationType.GROUP_NORM:
-      x = self.norm(x)
-      x = input_tensor.view(B, C, H * W)
+      x = self.norm(input_tensor)
+      x = x.view(B, C, H * W)
       x = x.transpose(-1, -2)
     else:
       x = input_tensor.view(B, C, H * W)
@@ -181,10 +180,9 @@ def forward(
     """
     residual = input_tensor
     B, C, H, W = input_tensor.shape
-    x = input_tensor
     if self.config.normalization_config.type == layers_cfg.NormalizationType.GROUP_NORM:
-      x = self.norm(x)
-      x = input_tensor.view(B, C, H * W)
+      x = self.norm(input_tensor)
+      x = x.view(B, C, H * W)
       x = x.transpose(-1, -2)
     else:
       x = input_tensor.view(B, C, H * W)
@@ -222,10 +220,9 @@ def __init__(
   def forward(self, input_tensor: torch.Tensor) -> torch.Tensor:
     residual = input_tensor
     B, C, H, W = input_tensor.shape
-    x = input_tensor
     if self.config.normalization_config.type == layers_cfg.NormalizationType.GROUP_NORM:
-      x = self.norm(x)
-      x = input_tensor.view(B, C, H * W)
+      x = self.norm(input_tensor)
+      x = x.view(B, C, H * W)
       x = x.transpose(-1, -2)
     else:
       x = input_tensor.view(B, C, H * W)