tf.keras.layers.Conv3DTranspose throws error when backprop
Same problem can be found in tf.keras.layers.LocallyConnected2D. Here is the repo code:
```
filters = 2048
kernel_size_0 = 3
kernel_size_1 = 3
kernel_size = [kernel_size_0,kernel_size_1,]
strides_0 = 1
strides_1 = 1
strides = [strides_0,strides_1,]
padding = "valid"
data_format = None
activation = "linear"
use_bias = True
kernel_initializer = None
bias_initializer = None
kernel_regularizer = None
bias_regularizer = None
activity_regularizer = None
kernel_constraint = None
bias_constraint = None
implementation = 1
__input___0_tensor = tf.random.uniform([1, 6, 10, 1], minval=1.1, maxval=10, dtype=tf.float64)
__input___0 = tf.identity(__input___0_tensor)
LocallyConnected2D_class = tf.keras.layers.LocallyConnected2D(filters, kernel_size, strides=strides, padding=padding, data_format=data_format, activation=activation, use_bias=use_bias, kernel_initializer=kernel_initializer, bias_initializer=bias_initializer, kernel_regularizer=kernel_regularizer, bias_regularizer=bias_regularizer, activity_regularizer=activity_regularizer, kernel_constraint=kernel_constraint, bias_constraint=bias_constraint, implementation=implementation)

layer = LocallyConnected2D_class
inputs = __input___0

with tf.GradientTape() as g:
    g.watch(inputs)
    res = layer(inputs)
print(res.shape)
grad = g.jacobian(res, inputs)
print(grad)
``` Hi @cheyennee ,

I have tested the given code on colab and its working fine.Please refer attached [gist](https://colab.sandbox.google.com/gist/SuryanarayanaY/ef3f074245957fa1ce975893ed897f39/62426.ipynb). 

Please note that I have reduced the no of filters due to Memory constraints but it should not affect the reported behaviour.

Could you please verify the behaviour attached. Can you confirm whether the issue with Windows Package as it will download intel package?

Thanks! Hi @SuryanarayanaY,
It appears that the number of filters does make a difference 😂. In your [gist](https://colab.research.google.com/gist/SuryanarayanaY/ef3f074245957fa1ce975893ed897f39/62426.ipynb#scrollTo=1chUD-vIX7Dg), where the number of filters is set to 40, there are no crashes. However, I repo above code in colab, in the provided [gist](https://colab.research.google.com/drive/1zOQm4FfPCepRdBH9M3-ef6KV8_wbF3jA), where the number of filters is increased to 1792, it crashes, and the error message suggests a potential OOM issue.  Same problem can be found in tf.keras.layers.UpSampling2D. Here is the repo code:
```
import tensorflow as tf

size_0 = 224
size_1 = 224
size = [size_0,size_1, ]
data_format = None
interpolation = "nearest"
__input___0_tensor = tf.random.uniform([1, 8, 8, 1], minval=0.0, maxval=2.84955773162074, dtype=tf.float64)
__input___0 = tf.identity(__input___0_tensor)
UpSampling2D_class = tf.keras.layers.UpSampling2D(size=size, data_format=data_format, interpolation=interpolation)


layer = UpSampling2D_class
inputs = __input___0

with tf.GradientTape() as g:
    g.watch(inputs)
    res = layer(inputs)
print(res.shape)
grad = g.jacobian(res, inputs)
print(grad)
```
output is:
```
(1, 1792, 1792, 1)
---------------------------------------------------------------------------
ResourceExhaustedError                    Traceback (most recent call last)
[<ipython-input-1-8abc242ca0fd>](https://localhost:8080/#) in <cell line: 20>()
     18     res = layer(inputs)
     19 print(res.shape)
---> 20 grad = g.jacobian(res, inputs)
     21 print(grad)

3 frames
[/usr/local/lib/python3.10/dist-packages/tensorflow/python/eager/execute.py](https://localhost:8080/#) in quick_execute(op_name, num_outputs, inputs, attrs, ctx, name)
     58         for t in inputs
     59     ]
---> 60     tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
     61                                         inputs, attrs, num_outputs)
     62   except core._NotOkStatusException as e:

ResourceExhaustedError: Graph execution error:

Detected at node gradient_tape/UnsortedSegmentSum/pfor/UnsortedSegmentSum defined at (most recent call last):
  File "/usr/lib/python3.10/runpy.py", line 196, in _run_module_as_main

  File "/usr/lib/python3.10/runpy.py", line 86, in _run_code

  File "/usr/local/lib/python3.10/dist-packages/colab_kernel_launcher.py", line 37, in <module>

  File "/usr/local/lib/python3.10/dist-packages/traitlets/config/application.py", line 992, in launch_instance

  File "/usr/local/lib/python3.10/dist-packages/ipykernel/kernelapp.py", line 619, in start

  File "/usr/local/lib/python3.10/dist-packages/tornado/platform/asyncio.py", line 195, in start

  File "/usr/lib/python3.10/asyncio/base_events.py", line 603, in run_forever

  File "/usr/lib/python3.10/asyncio/base_events.py", line 1909, in _run_once

  File "/usr/lib/python3.10/asyncio/events.py", line 80, in _run

  File "/usr/local/lib/python3.10/dist-packages/tornado/ioloop.py", line 685, in <lambda>

  File "/usr/local/lib/python3.10/dist-packages/tornado/ioloop.py", line 738, in _run_callback

  File "/usr/local/lib/python3.10/dist-packages/tornado/gen.py", line 825, in inner

  File "/usr/local/lib/python3.10/dist-packages/tornado/gen.py", line 786, in run

  File "/usr/local/lib/python3.10/dist-packages/ipykernel/kernelbase.py", line 361, in process_one

  File "/usr/local/lib/python3.10/dist-packages/tornado/gen.py", line 234, in wrapper

  File "/usr/local/lib/python3.10/dist-packages/ipykernel/kernelbase.py", line 261, in dispatch_shell

  File "/usr/local/lib/python3.10/dist-packages/tornado/gen.py", line 234, in wrapper

  File "/usr/local/lib/python3.10/dist-packages/ipykernel/kernelbase.py", line 539, in execute_request

  File "/usr/local/lib/python3.10/dist-packages/tornado/gen.py", line 234, in wrapper

  File "/usr/local/lib/python3.10/dist-packages/ipykernel/ipkernel.py", line 302, in do_execute

  File "/usr/local/lib/python3.10/dist-packages/ipykernel/zmqshell.py", line 539, in run_cell

  File "/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py", line 2975, in run_cell

  File "/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py", line 3030, in _run_cell

  File "/usr/local/lib/python3.10/dist-packages/IPython/core/async_helpers.py", line 78, in _pseudo_sync_runner

  File "/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py", line 3257, in run_cell_async

  File "/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py", line 3473, in run_ast_nodes

  File "/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py", line 3553, in run_code

  File "<ipython-input-1-8abc242ca0fd>", line 20, in <cell line: 20>

  File "/usr/local/lib/python3.10/dist-packages/tensorflow/python/ops/parallel_for/control_flow_ops.py", line 212, in f

OOM when allocating tensor with shape[10312216477696] and type float on /job:localhost/replica:0/task:0/device:GPU:0 by allocator GPU_0_bfc
	 [[{{node gradient_tape/UnsortedSegmentSum/pfor/UnsortedSegmentSum}}]]
Hint: If you want to see a list of allocated tensors when OOM happens, add report_tensor_allocations_upon_oom to RunOptions for current allocation info. This isn't available when running in Eager mode.
 [Op:__inference_f_162]
```
It may be that the size is too large and the memory is overflowing. Same problem can be found in tf.keras.layers.ZeroPadding3D. Here is the repo code:
```
padding = 16
data_format = None
__input___0_tensor = tf.random.uniform([1, 1, 2, 2, 3], minval=0.8510533546655319,maxval=3.0,dtype=tf.float64)
__input___0 = tf.identity(__input___0_tensor)
ZeroPadding3D_class = tf.keras.layers.ZeroPadding3D(padding=padding,data_format=data_format)

layer = ZeroPadding3D_class
inputs = __input___0

with tf.GradientTape() as g:
    g.watch(inputs)
    res = layer(inputs)
print(res.shape)
grad = g.jacobian(res, inputs)
print(grad)
``` > Hi @SuryanarayanaY, It appears that the number of filters does make a difference 😂. In your [gist](https://colab.research.google.com/gist/SuryanarayanaY/ef3f074245957fa1ce975893ed897f39/62426.ipynb#scrollTo=1chUD-vIX7Dg), where the number of filters is set to 40, there are no crashes. However, I repo above code in colab, in the provided [gist](https://colab.research.google.com/drive/1zOQm4FfPCepRdBH9M3-ef6KV8_wbF3jA), where the number of filters is increased to 1792, it crashes, and the error message suggests a potential OOM issue.

Hi @cheyennee , OOM errors are not problem with Tensorflow and with low number of filters it is indeed working.


Oh Sorry,I missed this from your logs.

```
OOM when allocating tensor with shape[50982027264] and type float on /job:localhost/replica:0/task:0/device:GPU:0 by allocator GPU_0_bfc
	 [[{{node gradient_tape/UnsortedSegmentSum/pfor/UnsortedSegmentSum}}]]
Hint: If you want to see a list of allocated tensors when OOM happens, add report_tensor_allocations_upon_oom to RunOptions for current allocation info. This isn't available when running in Eager mode.
 [Op:__inference_f_52731]
```
This is OOM error and not a problem with Tensorflow. With higher input size this is intended. Hi @SuryanarayanaY , I understand. It appears that this issue is related to the device I am currently using. I'll reduce the no of filters. Thank you! 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/62426">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/62426">No</a>

### Issue type

Bug

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

No

### Source

source

### TensorFlow version

tf 2.14.0

### Custom code

Yes

### OS platform and distribution

windows colab

### Mobile device

_No response_

### Python version

3.10

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current behavior?

During backpropagation, the API crashes. Based on the error message, it appears to be indicative of an OOM situation.

### Standalone code to reproduce the issue

```shell
filters = 1792
kernel_size_0 = 3
kernel_size_1 = 3
kernel_size_2 = 3
kernel_size = [kernel_size_0,kernel_size_1,kernel_size_2,]
strides_0 = 1
strides_1 = 1
strides_2 = 1
strides = [strides_0,strides_1,strides_2,]
padding = "valid"
output_padding = None
data_format = "channels_last"
dilation_rate_0 = 1
dilation_rate_1 = 1
dilation_rate_2 = 1
dilation_rate = [dilation_rate_0,dilation_rate_1,dilation_rate_2,]
activation = "linear"
use_bias = True
kernel_initializer = None
bias_initializer = None
kernel_regularizer = None
bias_regularizer = None
activity_regularizer = None
kernel_constraint = None
bias_constraint = None
__input___0_tensor = tf.random.uniform([2, 5, 1, 1, 1], minval=0, maxval=1, dtype=tf.float32)
__input___0 = tf.identity(__input___0_tensor)
Conv3DTranspose_class = tf.keras.layers.Conv3DTranspose(filters, kernel_size, strides=strides, padding=padding, output_padding=output_padding, data_format=data_format, dilation_rate=dilation_rate, activation=activation, use_bias=use_bias, kernel_initializer=kernel_initializer, bias_initializer=bias_initializer, kernel_regularizer=kernel_regularizer, bias_regularizer=bias_regularizer, activity_regularizer=activity_regularizer, kernel_constraint=kernel_constraint, bias_constraint=bias_constraint)

layer = Conv3DTranspose_class
inputs = __input___0

with tf.GradientTape() as g:
    g.watch(inputs)
    res = layer(inputs)
print(res.shape)
grad = g.jacobian(res, inputs) # Error
```


### Relevant log output

```shell
(2, 7, 3, 3, 1792)
---------------------------------------------------------------------------
ResourceExhaustedError                    Traceback (most recent call last)
<ipython-input-6-15c96e537b1f> in <cell line: 47>()
     45     res = layer(inputs)
     46 print(res.shape) # (1, 5, 6, 2)
---> 47 grad = g.jacobian(res, inputs) # Error

3 frames
/usr/local/lib/python3.10/dist-packages/tensorflow/python/eager/execute.py in quick_execute(op_name, num_outputs, inputs, attrs, ctx, name)
     58         for t in inputs
     59     ]
---> 60     tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
     61                                         inputs, attrs, num_outputs)
     62   except core._NotOkStatusException as e:

ResourceExhaustedError: Graph execution error:

Detected at node gradient_tape/UnsortedSegmentSum/pfor/UnsortedSegmentSum defined at (most recent call last):
  File "/usr/lib/python3.10/runpy.py", line 196, in _run_module_as_main

  File "/usr/lib/python3.10/runpy.py", line 86, in _run_code

  File "/usr/local/lib/python3.10/dist-packages/colab_kernel_launcher.py", line 37, in <module>

  File "/usr/local/lib/python3.10/dist-packages/traitlets/config/application.py", line 992, in launch_instance

  File "/usr/local/lib/python3.10/dist-packages/ipykernel/kernelapp.py", line 619, in start

  File "/usr/local/lib/python3.10/dist-packages/tornado/platform/asyncio.py", line 195, in start

  File "/usr/lib/python3.10/asyncio/base_events.py", line 603, in run_forever

  File "/usr/lib/python3.10/asyncio/base_events.py", line 1909, in _run_once

  File "/usr/lib/python3.10/asyncio/events.py", line 80, in _run

  File "/usr/local/lib/python3.10/dist-packages/tornado/ioloop.py", line 685, in <lambda>

  File "/usr/local/lib/python3.10/dist-packages/tornado/ioloop.py", line 738, in _run_callback

  File "/usr/local/lib/python3.10/dist-packages/tornado/gen.py", line 825, in inner

  File "/usr/local/lib/python3.10/dist-packages/tornado/gen.py", line 786, in run

  File "/usr/local/lib/python3.10/dist-packages/ipykernel/kernelbase.py", line 361, in process_one

  File "/usr/local/lib/python3.10/dist-packages/tornado/gen.py", line 234, in wrapper

  File "/usr/local/lib/python3.10/dist-packages/ipykernel/kernelbase.py", line 261, in dispatch_shell

  File "/usr/local/lib/python3.10/dist-packages/tornado/gen.py", line 234, in wrapper

  File "/usr/local/lib/python3.10/dist-packages/ipykernel/kernelbase.py", line 539, in execute_request

  File "/usr/local/lib/python3.10/dist-packages/tornado/gen.py", line 234, in wrapper

  File "/usr/local/lib/python3.10/dist-packages/ipykernel/ipkernel.py", line 302, in do_execute

  File "/usr/local/lib/python3.10/dist-packages/ipykernel/zmqshell.py", line 539, in run_cell

  File "/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py", line 2975, in run_cell

  File "/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py", line 3030, in _run_cell

  File "/usr/local/lib/python3.10/dist-packages/IPython/core/async_helpers.py", line 78, in _pseudo_sync_runner

  File "/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py", line 3257, in run_cell_async

  File "/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py", line 3473, in run_ast_nodes

  File "/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py", line 3553, in run_code

  File "<ipython-input-6-15c96e537b1f>", line 47, in <cell line: 47>

  File "/usr/local/lib/python3.10/dist-packages/tensorflow/python/ops/parallel_for/control_flow_ops.py", line 212, in f

OOM when allocating tensor with shape[50982027264] and type float on /job:localhost/replica:0/task:0/device:GPU:0 by allocator GPU_0_bfc
	 [[{{node gradient_tape/UnsortedSegmentSum/pfor/UnsortedSegmentSum}}]]
Hint: If you want to see a list of allocated tensors when OOM happens, add report_tensor_allocations_upon_oom to RunOptions for current allocation info. This isn't available when running in Eager mode.
 [Op:__inference_f_52731]
```
