Fix yapf

runame · Dec 11, 2023 · 03a200f · 03a200f
1 parent dedef80
commit 03a200f
Showing 1 changed file with 2 additions and 6 deletions.
diff --git a/algorithmic_efficiency/workloads/wmt/wmt_pytorch/models.py b/algorithmic_efficiency/workloads/wmt/wmt_pytorch/models.py
@@ -229,9 +229,7 @@ def __init__(self,
     self.enable_nested_tensor = enable_nested_tensor
     self.mask_check = mask_check
 
-  def forward(self,
-              src: Tensor,
-              mask: Optional[Tensor] = None) -> Tensor:
+  def forward(self, src: Tensor, mask: Optional[Tensor] = None) -> Tensor:
     """Pass the input through the encoder layers in turn.
 
     Args:
@@ -494,9 +492,7 @@ def __init__(self,
 
     self.activation = activation
 
-  def forward(self,
-              src: Tensor,
-              src_mask: Optional[Tensor] = None) -> Tensor:
+  def forward(self, src: Tensor, src_mask: Optional[Tensor] = None) -> Tensor:
     r"""Pass the input through the encoder layer.
 
     Args: