make style

huggingface · Jul 27, 2022 · 8238406 · 8238406
1 parent 6ff64d1
commit 8238406
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 5 deletions.
diff --git a/examples/flax/text-classification/run_flax_glue.py b/examples/flax/text-classification/run_flax_glue.py
@@ -36,7 +36,7 @@
 import optax
 import transformers
 from flax import struct, traverse_util
-from flax.jax_utils import replicate, unreplicate, pad_shard_unpad
+from flax.jax_utils import pad_shard_unpad, replicate, unreplicate
 from flax.training import train_state
 from flax.training.common_utils import get_metrics, onehot, shard
 from huggingface_hub import Repository
@@ -622,7 +622,9 @@ def eval_step(state, batch):
                     position=2,
                 ):
                     labels = batch.pop("labels")
-                    predictions = pad_shard_unpad(p_eval_step)(state, batch, min_device_batch=per_device_eval_batch_size)
+                    predictions = pad_shard_unpad(p_eval_step)(
+                        state, batch, min_device_batch=per_device_eval_batch_size
+                    )
                     metric.add_batch(predictions=np.array(predictions), references=labels)
 
                 eval_metric = metric.compute()

diff --git a/examples/flax/token-classification/run_flax_ner.py b/examples/flax/token-classification/run_flax_ner.py
@@ -37,7 +37,7 @@
 import optax
 import transformers
 from flax import struct, traverse_util
-from flax.jax_utils import replicate, unreplicate, pad_shard_unpad
+from flax.jax_utils import pad_shard_unpad, replicate, unreplicate
 from flax.training import train_state
 from flax.training.common_utils import get_metrics, onehot, shard
 from huggingface_hub import Repository

diff --git a/examples/flax/vision/run_image_classification.py b/examples/flax/vision/run_image_classification.py
@@ -40,7 +40,7 @@
 import optax
 import transformers
 from flax import jax_utils
-from flax.jax_utils import unreplicate, pad_shard_unpad
+from flax.jax_utils import pad_shard_unpad, unreplicate
 from flax.training import train_state
 from flax.training.common_utils import get_metrics, onehot, shard, shard_prng_key
 from huggingface_hub import Repository
@@ -534,7 +534,9 @@ def eval_step(params, batch):
         eval_step_progress_bar = tqdm(total=eval_steps, desc="Evaluating...", position=2, leave=False)
         for batch in eval_loader:
             # Model forward
-            metrics = pad_shard_unpad(p_eval_step, static_return=True)(state.params, batch, min_device_batch=per_device_eval_batch_size)
+            metrics = pad_shard_unpad(p_eval_step, static_return=True)(
+                state.params, batch, min_device_batch=per_device_eval_batch_size
+            )
             eval_metrics.append(metrics)
 
             eval_step_progress_bar.update(1)