[Flax] Correct logging steps flax (#12515)

* fix_torch_device_generate_test

* remove @

* push

examples/flax/language-modeling/run_clm_flax.py

@@ -574,7 +574,7 @@ def main():
 
             cur_step = epoch * (len(train_dataset) // train_batch_size) + step
 
-            if cur_step % training_args.logging_steps and cur_step > 0:
+            if cur_step % training_args.logging_steps == 0 and cur_step > 0:
                 # Save metrics
                 train_metric = unreplicate(train_metric)
                 train_time += time.time() - train_start

examples/flax/language-modeling/run_mlm_flax.py

@@ -608,7 +608,7 @@ if __name__ == "__main__":
 
             cur_step = epoch * num_train_samples + step
 
-            if cur_step % training_args.logging_steps and cur_step > 0:
+            if cur_step % training_args.logging_steps == 0 and cur_step > 0:
                 # Save metrics
                 train_metric = jax_utils.unreplicate(train_metric)
                 train_time += time.time() - train_start

examples/flax/language-modeling/run_t5_mlm_flax.py

@@ -724,7 +724,7 @@ if __name__ == "__main__":
 
             cur_step = epoch * num_train_samples + step
 
-            if cur_step % training_args.logging_steps and cur_step > 0:
+            if cur_step % training_args.logging_steps == 0 and cur_step > 0:
                 # Save metrics
                 train_metric = jax_utils.unreplicate(train_metric)
                 train_time += time.time() - train_start