lint

Author: entrpn

src/maxdiffusion/models/attention_flax.py

@@ -193,15 +193,11 @@ def cudnn_flash_attention(
     value = nn.with_logical_constraint(value, axis_names)
 
     @functools.partial(
-      shard_map.shard_map,
-      mesh=self.mesh,
-      in_specs=(
-        axis_names,
-        axis_names,
-        axis_names
-      ),
-      out_specs=axis_names,
-      check_rep=False
+        shard_map.shard_map,
+        mesh=self.mesh,
+        in_specs=(axis_names, axis_names, axis_names),
+        out_specs=axis_names,
+        check_rep=False,
     )
     def wrap_flash_attention(query, key, value):
       return jax.vmap(self.dpa_layer)(query, key, value, mask=None)

src/maxdiffusion/train_sdxl.py

@@ -49,6 +49,7 @@ def main(argv: Sequence[str]) -> None:
   import os
   import tensorflow as tf
   import torch
+
   tf.config.set_visible_devices([], "GPU")
   os.environ["TF_CPP_MIN_LOG_LEVEL"] = "0"
   torch.set_default_device("cpu")