lint

mlcommons · Dec 8, 2023 · 95e85bf · 95e85bf
1 parent a00a031
commit 95e85bf
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/algorithmic_efficiency/workloads/fastmri/fastmri_jax/models.py b/algorithmic_efficiency/workloads/fastmri/fastmri_jax/models.py
@@ -126,7 +126,9 @@ def __call__(self, x, train=True):
       output = jnp.concatenate((output, downsample_layer), axis=-1)
       output = conv(output, train)
 
-    output = nn.Conv(self.out_channels, kernel_size=(1, 1), strides=(1, 1))(output)
+    output = nn.Conv(
+        self.out_channels, kernel_size=(1, 1), strides=(1, 1))(
+            output)
     return output.squeeze(-1)