Can the resnet model written in the tf_slim library call the MirroredStrategy strategy to achieve data parallel training?
Hi @yeoways ,

To use Distribution strategies first we need to create a strategy from applicable Distribution strategy class and then build the model inside that strategy scope. Then remaining steps are as usual.

```
strategy = tf.distribute.MirroredStrategy()

with strategy.scope():
  # Model building/compiling need to be within `strategy.scope()`.
  multi_GPU_model = <call_your_model()>
```

You can refer the attached tutorials [1](https://www.tensorflow.org/tutorials/distribute/keras) ,[2](https://www.tensorflow.org/tutorials/distribute/multi_worker_with_keras) and [3](https://www.tensorflow.org/tutorials/distribute/multi_worker_with_ctl)  for going in details on how to implement distribution strategies.


The attached code seems related to TF1.x versions. You can refer the resource 3 attached above for distribution training with custom training loop where we can override the `train_step` function to define our own training steps. 
 This issue is stale because it has been open for 7 days with no activity. It will be closed if no further activity occurs. Thank you. This issue was closed because it has been inactive for 7 days since being marked as stale. Please reopen if you'd like to work on this further. Are you satisfied with the resolution of your issue?
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=Yes&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/61239">Yes</a>
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=No&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/61239">No</a>

### Issue type

Support

### Have you reproduced the bug with TensorFlow Nightly?

Yes

### Source

source

### TensorFlow version

tf 2.11

### Custom code

Yes

### OS platform and distribution

Linux Ubuntu 16.04

### Mobile device

Linux Ubuntu 16.04

### Python version

3.9

### Bazel version

5.1.1

### GCC/compiler version

9.4

### CUDA/cuDNN version

_No response_

### GPU model and memory

Tesla P100 12GB

### Current behavior?


The distributed_train_step method is only called once in the first epoch, and the rest are not called

### Standalone code to reproduce the issue

```shell
from tensorflow.keras.datasets import cifar10
from tensorflow.keras.utils import to_categorical
from tensorflow.keras.optimizers import Adam
(x_train, y_train), (x_test, y_test) = cifar10.load_data()
x_train = x_train.astype('float32') / 255.0
x_test = x_test.astype('float32') / 255.0
y_train = to_categorical(y_train, num_classes=10)
y_test = to_categorical(y_test, num_classes=10)
# 输入占位符
inputs = tf.placeholder(tf.float32, shape=(None, 32, 32, 3))
labels = tf.placeholder(tf.float32, shape=(None, 10))
batch_size = 32
num_epochs = 10
num_batches = len(x_train) // batch_size

strategy = tf2.distribute.MirroredStrategy(devices=["GPU:0", "GPU:1","GPU:2", "GPU:3"])
# 将训练数据集分发到多个GPU上
train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train)).shuffle(len(x_train)).batch(batch_size)
dist_train_dataset = strategy.experimental_distribute_dataset(train_dataset)
with strategy.scope():
    # 构建ResNet模型
    net, end_points = resnet_v2_152(inputs,10)
    net = tf.squeeze(net, axis=[1, 2])  # 移除维度为 1 的高度和宽度维度

    # 定义损失函数和优化器
    optimizer = tf.train.AdamOptimizer(learning_rate=0.001)


@tf.function
def train_step(input):
    x, y = input
    print('x_shape:',x.shape)
    print('y_shape:',y.shape)
    with tf.GradientTape() as tape:
        logits,_ = resnet_v2_152(x,10,reuse = True)
        print('logits:',logits.shape)
        logits = tf.squeeze(logits, axis=[1, 2])  # 移除维度为 1 的高度和宽度维度
        print('logits_squeeze:',logits.shape)
        loss_value = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(labels=y, logits=logits))
        print('loss_value:',loss_value.shape)
    grads = tape.gradient(loss_value, tf.trainable_variables())
    optimizer.apply_gradients(zip(grads, tf.trainable_variables()))
    correct_predictions = tf.equal(tf.argmax(logits, axis=1), tf.argmax(y, axis=1))
    accuracy = tf.reduce_mean(tf.cast(correct_predictions, tf.float32))
    return loss_value, accuracy
@tf.function
def distributed_train_step(dataset_inputs):
        total_loss = 0.0
        total_acc = 0.0
        num_batches = 0
        print("distributed_train_step function start")
        print("dataset_inputs:{}".format(dataset_inputs))
        for x in dataset_inputs:
            print("x:{}".format(x))
            per_replica_losses, per_replica_accuracies = strategy.run(train_step, args=(x,))
            print("per_replica_losses:{},per_replica_acc:{}".format(per_replica_losses,per_replica_accuracies))
            total_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses, axis=None)
            total_acc = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_accuracies, axis=None)
            num_batches +=1
         return total_loss / tf.cast(num_batches, dtype=tf.float32),total_acc / tf.cast(num_batches, dtype=tf.float32)
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    for epoch in range(num_epochs):
        start_time = time.time()
        train_loss,train_acc = distributed_train_step(dist_train_dataset)
        template = ("Epoch {}, Loss: {}, Accuracy: {}")
        print(template.format(epoch + 1, train_loss, train_acc))
        print("sess epoch:{},time:{}".format(epoch+1,time.time()-start_time))
        print()
```


### Relevant log output

_No response_