fit() fails with CUDNN_STATUS_BAD_PARAM when using Conv3D and multi-GPU MirroredStrategy
Hi @ohinds ,

For distribution training across multiple devices on same server you need to keep `batch_szie = 1* no of GPU devices` incase you want to calculate gradients after each single batch. Since you have 4 GPUs you need to keep batch_size=1*4 as minimum (multiples of 4). The no of devices can be found by using `mirrored_strategy.num_replicas_in_sync` for your case.

Hence you need to pass `batch_size=batch_size*mirrored_strategy.num_replicas_in_sync`

Please try the above change and let us know if it works. Thanks!

 Hi @SuryanarayanaY,
Thanks! Changing the fit line to
```
model.fit(x_data, y_data, batch_size=mirrored_strategy.num_replicas_in_sync, epochs=1, verbose=1)
```

fixed the issue.  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/60935">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/60935">No</a>

<details><summary>Click to expand!</summary> 
 
 ### Issue Type

Bug

### Have you reproduced the bug with TF nightly?

Yes

### Source

binary

### Tensorflow Version

v1.12.1-95675-g47602c0bad8 2.14.0-dev20230620

### Custom Code

Yes

### OS Platform and Distribution

Rocky Linux release 8.6 (Green Obsidian)

### Mobile device

_No response_

### Python version

3.8.10

### Bazel version

_No response_

### GCC/Compiler version

_No response_

### CUDA/cuDNN version

cuda_11.8.r11.8/compiler.31833905_0 / cuDNN version 8600

### GPU model and memory

4 NVIDIA A100s w/ 80GB each

### Current Behaviour?

When executing a model fit that includes a `Conv3D` layer on multiple GPUs, I'm encountering a `CUDNN_STATUS_BAD_PARAM` error in the gradient computation step. No errors occur when running on a single GPU, nor when I swap out `Conv3D` with `AveragePooling3D` or `Conv2D`. However, `Conv3DTranspose` also fails.

```none
  (0) UNKNOWN:  CUDNN_STATUS_BAD_PARAM
in tensorflow/compiler/xla/stream_executor/cuda/cuda_dnn.cc(3549): 'tensor' CUDNN_BACKEND_TENSOR_DESCRIPTOR: Check and Set the CUDNN_ATTR_TENSOR_DIMENSIONS Correctly
         [[{{node gradient_tape/replica_2/model/conv3d/Conv3D/Conv3DBackpropFilterV2}}]]
         [[div_no_nan/ReadVariableOp_1/_52]]
         [[group_deps/_95]]
         [[Adam/update_2_2/AssignAddVariableOp/_119]]
         [[group_deps/_103]]
```

With the `Graph execution error` traceback:
```none
Traceback (most recent call last):
  File "conv3_multi_gpu_fail_repro.py", line 25, in <module>
    model.fit(x_data, y_data, batch_size=1, epochs=1, verbose=1)
  File "/***/env/lib/python3.8/site-packages/keras/src/utils/traceback_utils.py", line 70, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/***/env/lib/python3.8/site-packages/tensorflow/python/eager/execute.py", line 53, in quick_execute
    tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
tensorflow.python.framework.errors_impl.UnknownError: Graph execution error:

Detected at node gradient_tape/replica_3/model/conv3d/Conv3D/Conv3DBackpropFilterV2 defined at (most recent call last):
  File "/usr/lib/python3.8/threading.py", line 890, in _bootstrap
    self._bootstrap_inner()
  File "/usr/lib/python3.8/threading.py", line 890, in _bootstrap
    self._bootstrap_inner()
  File "/usr/lib/python3.8/threading.py", line 932, in _bootstrap_inner
    self.run()
  File "/usr/lib/python3.8/threading.py", line 890, in _bootstrap
    self._bootstrap_inner()
  File "/usr/lib/python3.8/threading.py", line 932, in _bootstrap_inner
    self.run()
  File "/***/env/lib/python3.8/site-packages/keras/src/engine/training.py", line 1348, in run_step
    outputs = model.train_step(data)
  File "/usr/lib/python3.8/threading.py", line 890, in _bootstrap
    self._bootstrap_inner()
  File "/usr/lib/python3.8/threading.py", line 932, in _bootstrap_inner
    self.run()
  File "/***/env/lib/python3.8/site-packages/keras/src/engine/training.py", line 1348, in run_step
    outputs = model.train_step(data)
  File "/***/env/lib/python3.8/site-packages/keras/src/engine/training.py", line 1129, in train_step
    self.optimizer.minimize(loss, self.trainable_variables, tape=tape)
  File "/usr/lib/python3.8/threading.py", line 890, in _bootstrap
    self._bootstrap_inner()
  File "/usr/lib/python3.8/threading.py", line 932, in _bootstrap_inner
    self.run()
  File "/***/env/lib/python3.8/site-packages/keras/src/engine/training.py", line 1348, in run_step
    outputs = model.train_step(data)
  File "/***/env/lib/python3.8/site-packages/keras/src/engine/training.py", line 1129, in train_step
    self.optimizer.minimize(loss, self.trainable_variables, tape=tape)
  File "/***/env/lib/python3.8/site-packages/keras/src/optimizers/optimizer.py", line 543, in minimize
    grads_and_vars = self.compute_gradients(loss, var_list, tape)
  File "/usr/lib/python3.8/threading.py", line 890, in _bootstrap
    self._bootstrap_inner()
  File "/usr/lib/python3.8/threading.py", line 932, in _bootstrap_inner
    self.run()
  File "/***/env/lib/python3.8/site-packages/keras/src/engine/training.py", line 1348, in run_step
    outputs = model.train_step(data)
  File "/***/env/lib/python3.8/site-packages/keras/src/engine/training.py", line 1129, in train_step
    self.optimizer.minimize(loss, self.trainable_variables, tape=tape)
  File "/***/env/lib/python3.8/site-packages/keras/src/optimizers/optimizer.py", line 543, in minimize
    grads_and_vars = self.compute_gradients(loss, var_list, tape)
  File "/***/env/lib/python3.8/site-packages/keras/src/optimizers/optimizer.py", line 276, in compute_gradients
    grads = tape.gradient(loss, var_list)
```

I'm running from the `tensorflow/tensorflow:nightly-gpu` docker image.


### Standalone code to reproduce the issue

```python
import tensorflow as tf
from tensorflow.keras import layers, models

input_shape = (28, 28, 28, 1)
num_samples = 10

x_data = tf.random.uniform((num_samples, *input_shape), 0, 1)
y_data = tf.random.uniform((num_samples, *input_shape), 0, 1)

multi_gpu=True # <== fails
#multi_gpu=False # <== works
devices = [] if multi_gpu else ['/gpu:0']

mirrored_strategy = tf.distribute.MirroredStrategy(devices=devices)
print(f"{mirrored_strategy.num_replicas_in_sync} replica(s)")

with mirrored_strategy.scope():
    inputs = layers.Input(shape=input_shape)
    outputs = layers.Conv3D(1, 1)(inputs) # <== fails
    #outputs = layers.AveragePooling3D(1)(inputs) # <== works
    #outputs = layers.Conv2D(1, 1)(inputs) # <== works
    model = models.Model(inputs=inputs, outputs=outputs)
    model.compile(optimizer='adam', loss='binary_crossentropy')

model.fit(x_data, y_data, batch_size=1, epochs=1, verbose=1)
```


### Relevant log output

_No response_</details>