update multi gpu

bubbliiiing · Apr 16, 2022 · 1b86e78 · 1b86e78
1 parent 2d10e3c
commit 1b86e78
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/utils/utils_fit.py b/utils/utils_fit.py
@@ -32,7 +32,7 @@ def train_step(images, labels, net, optimizer, loss, aux_branch, metrics):
         @tf.function
         def distributed_train_step(images, labels, net, optimizer, loss, aux_branch, metrics):
             per_replica_losses, per_replica_score = strategy.run(train_step, args=(images, labels, net, optimizer, loss, aux_branch, metrics))
-            return strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_losses,axis=None), strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_score, axis=None)
+            return strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_losses, axis=None), strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_score, axis=None)
         return distributed_train_step
 
 @tf.function