TensorFlow distributed training works for at most 2 GPUs
@aeave If you want to use multiple GPUs, you have to manually specify what tensors to put on each GPU.
Could you please refer to [this](https://www.tensorflow.org/guide/distributed_training) document which explains in detail about the distribution strategy and multi-gpu. Thank you! Hello! Thank you for responding. I followed the steps at the following link when writing the code example: https://www.tensorflow.org/guide/distributed_training#use_tfdistributestrategy_with_keras_modelfit

I was thinking that the `tf.distribute.MirroredStrategy` API would automatically divide each batch evenly across multiple GPUs. This is the snippet of documentation from the TensorFlow website that led me to believe this:
![image](https://github.com/tensorflow/tensorflow/assets/77916424/5fb24e8d-6603-4b4c-bcb3-c7d30153b701)
 
When running the code with two GPUs passed into the `tf.distribute.MirroredStrategy()` constructor, I am able to monitor utilization of both of them during the model's training. During training, the two GPUs that I have selected increase in utilization, so the code seems to be working as intended for the case with only two GPUs. However, the code will error if I ask it to use more than two GPUs. If I do not pass in the `gpus` list at all into the `tf.distribute.MirroredStrategy()` constructor, I expect the code to automatically detect all 8 GPUs on the system. `strategy.num_replicas_in_sync` returns 8 in this case, so I assume it is working. However, the same error as with the case of passing only 3 GPUs into the constructor appears in the case with 8 GPUs as well. Thank you for taking the time to help me out! Hi @aeave ,

When using distribution strategy we need to create and compile the model in the strategy scope itself. Also try adding the argument `drop_remainder=True` to the` batch()` method . So please change code like below.

```
import tensorflow as tf
gpus = ["/gpu:0", "/gpu:1", "/gpu:2"]
strategy = tf.distribute.MirroredStrategy(gpus)
with strategy.scope():
    model = tf.keras.Sequential([tf.keras.layers.Dense(1, input_shape=(1,))])
    model.compile(loss="mse", optimizer="sgd")
dataset = tf.data.Dataset.from_tensors(([1.], [1.])).repeat(100).batch(16,drop_remainder=True)
model.fit(dataset)
```
Please execute the above code and let us know if it works.

Thanks! 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/61306">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/61306">No</a>

### Issue type

Bug

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

No

### Source

binary

### TensorFlow version

2.12.0

### Custom code

Yes

### OS platform and distribution

Linux Ubuntu 22.04

### Mobile device

_No response_

### Python version

3.9.16

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

CUDA 11.8, cuDNN 8.6.0.163

### GPU model and memory

_No response_

### Current behavior?

I have 8 NVIDIA H100s on this system. Below is the output of `nvidia-smi`.
![image](https://github.com/tensorflow/tensorflow/assets/77916424/7d027250-8312-4503-bf98-14b6d2aa0db9)

The following code works as intended if `gpus` includes at most 2 GPUs.
![image](https://github.com/tensorflow/tensorflow/assets/77916424/47f78e02-e92e-47cc-814b-4ddcba433a8b)

In the "Relevant log output" section, I have the error log from running the code with 3 GPUs in the `gpus` list.

Please let me know if you need any additional information!

### Standalone code to reproduce the issue

```shell
import tensorflow as tf
gpus = ["/gpu:0", "/gpu:1", "/gpu:2"]
strategy = tf.distribute.MirroredStrategy(gpus)
with strategy.scope():
    model = tf.keras.Sequential([tf.keras.layers.Dense(1, input_shape=(1,))])
model.compile(loss="mse", optimizer="sgd")
dataset = tf.data.Dataset.from_tensors(([1.], [1.])).repeat(100).batch(16)
model.fit(dataset)
```


### Relevant log output

```shell
2023-07-17 23:28:31.469519: I tensorflow/core/util/port.cc:110] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
2023-07-17 23:28:31.516738: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 AVX512F AVX512_VNNI AVX512_BF16 AVX_VNNI AMX_TILE AMX_INT8 AMX_BF16 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
2023-07-17 23:28:31.972313: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
2023-07-17 23:28:35.178825: W tensorflow/core/common_runtime/gpu/gpu_device.cc:2048] TensorFlow was not built with CUDA kernel binaries compatible with compute capability 9.0. CUDA kernels will be jit-compiled from PTX, which could take 30 minutes or longer.
2023-07-17 23:28:35.181183: W tensorflow/core/common_runtime/gpu/gpu_device.cc:2048] TensorFlow was not built with CUDA kernel binaries compatible with compute capability 9.0. CUDA kernels will be jit-compiled from PTX, which could take 30 minutes or longer.
2023-07-17 23:28:35.183528: W tensorflow/core/common_runtime/gpu/gpu_device.cc:2048] TensorFlow was not built with CUDA kernel binaries compatible with compute capability 9.0. CUDA kernels will be jit-compiled from PTX, which could take 30 minutes or longer.
2023-07-17 23:28:35.185905: W tensorflow/core/common_runtime/gpu/gpu_device.cc:2048] TensorFlow was not built with CUDA kernel binaries compatible with compute capability 9.0. CUDA kernels will be jit-compiled from PTX, which could take 30 minutes or longer.
2023-07-17 23:28:35.188234: W tensorflow/core/common_runtime/gpu/gpu_device.cc:2048] TensorFlow was not built with CUDA kernel binaries compatible with compute capability 9.0. CUDA kernels will be jit-compiled from PTX, which could take 30 minutes or longer.
2023-07-17 23:28:35.190559: W tensorflow/core/common_runtime/gpu/gpu_device.cc:2048] TensorFlow was not built with CUDA kernel binaries compatible with compute capability 9.0. CUDA kernels will be jit-compiled from PTX, which could take 30 minutes or longer.
2023-07-17 23:28:35.192895: W tensorflow/core/common_runtime/gpu/gpu_device.cc:2048] TensorFlow was not built with CUDA kernel binaries compatible with compute capability 9.0. CUDA kernels will be jit-compiled from PTX, which could take 30 minutes or longer.
2023-07-17 23:28:35.195198: W tensorflow/core/common_runtime/gpu/gpu_device.cc:2048] TensorFlow was not built with CUDA kernel binaries compatible with compute capability 9.0. CUDA kernels will be jit-compiled from PTX, which could take 30 minutes or longer.
2023-07-17 23:28:36.609484: W tensorflow/core/common_runtime/gpu/gpu_device.cc:2048] TensorFlow was not built with CUDA kernel binaries compatible with compute capability 9.0. CUDA kernels will be jit-compiled from PTX, which could take 30 minutes or longer.
2023-07-17 23:28:36.610685: W tensorflow/core/common_runtime/gpu/gpu_device.cc:2048] TensorFlow was not built with CUDA kernel binaries compatible with compute capability 9.0. CUDA kernels will be jit-compiled from PTX, which could take 30 minutes or longer.
2023-07-17 23:28:36.611908: W tensorflow/core/common_runtime/gpu/gpu_device.cc:2048] TensorFlow was not built with CUDA kernel binaries compatible with compute capability 9.0. CUDA kernels will be jit-compiled from PTX, which could take 30 minutes or longer.
2023-07-17 23:28:36.613166: W tensorflow/core/common_runtime/gpu/gpu_device.cc:2048] TensorFlow was not built with CUDA kernel binaries compatible with compute capability 9.0. CUDA kernels will be jit-compiled from PTX, which could take 30 minutes or longer.
2023-07-17 23:28:36.614370: W tensorflow/core/common_runtime/gpu/gpu_device.cc:2048] TensorFlow was not built with CUDA kernel binaries compatible with compute capability 9.0. CUDA kernels will be jit-compiled from PTX, which could take 30 minutes or longer.
2023-07-17 23:28:36.615584: W tensorflow/core/common_runtime/gpu/gpu_device.cc:2048] TensorFlow was not built with CUDA kernel binaries compatible with compute capability 9.0. CUDA kernels will be jit-compiled from PTX, which could take 30 minutes or longer.
2023-07-17 23:28:36.616799: W tensorflow/core/common_runtime/gpu/gpu_device.cc:2048] TensorFlow was not built with CUDA kernel binaries compatible with compute capability 9.0. CUDA kernels will be jit-compiled from PTX, which could take 30 minutes or longer.
2023-07-17 23:28:36.617975: W tensorflow/core/common_runtime/gpu/gpu_device.cc:2048] TensorFlow was not built with CUDA kernel binaries compatible with compute capability 9.0. CUDA kernels will be jit-compiled from PTX, which could take 30 minutes or longer.
2023-07-17 23:28:36.680264: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1635] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 78938 MB memory:  -> device: 0, name: NVIDIA H100 80GB HBM3, pci bus id: 0000:19:00.0, compute capability: 9.0
2023-07-17 23:28:36.682112: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1635] Created device /job:localhost/replica:0/task:0/device:GPU:1 with 78938 MB memory:  -> device: 1, name: NVIDIA H100 80GB HBM3, pci bus id: 0000:3b:00.0, compute capability: 9.0
2023-07-17 23:28:36.683772: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1635] Created device /job:localhost/replica:0/task:0/device:GPU:2 with 78938 MB memory:  -> device: 2, name: NVIDIA H100 80GB HBM3, pci bus id: 0000:4c:00.0, compute capability: 9.0
2023-07-17 23:28:36.685499: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1635] Created device /job:localhost/replica:0/task:0/device:GPU:3 with 78938 MB memory:  -> device: 3, name: NVIDIA H100 80GB HBM3, pci bus id: 0000:5d:00.0, compute capability: 9.0
2023-07-17 23:28:36.687179: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1635] Created device /job:localhost/replica:0/task:0/device:GPU:4 with 78938 MB memory:  -> device: 4, name: NVIDIA H100 80GB HBM3, pci bus id: 0000:9b:00.0, compute capability: 9.0
2023-07-17 23:28:36.688884: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1635] Created device /job:localhost/replica:0/task:0/device:GPU:5 with 78938 MB memory:  -> device: 5, name: NVIDIA H100 80GB HBM3, pci bus id: 0000:bb:00.0, compute capability: 9.0
2023-07-17 23:28:36.690581: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1635] Created device /job:localhost/replica:0/task:0/device:GPU:6 with 78938 MB memory:  -> device: 6, name: NVIDIA H100 80GB HBM3, pci bus id: 0000:cb:00.0, compute capability: 9.0
2023-07-17 23:28:36.692292: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1635] Created device /job:localhost/replica:0/task:0/device:GPU:7 with 78938 MB memory:  -> device: 7, name: NVIDIA H100 80GB HBM3, pci bus id: 0000:db:00.0, compute capability: 9.0
2023-07-17 23:28:39.806560: I tensorflow/core/common_runtime/executor.cc:1197] [/device:CPU:0] (DEBUG INFO) Executor start aborting (this does not indicate an error and you can ignore this message): INVALID_ARGUMENT: You must feed a value for placeholder tensor 'Placeholder/_1' with dtype float and shape [1]
         [[{{node Placeholder/_1}}]]
2023-07-17 23:28:39.806723: I tensorflow/core/common_runtime/executor.cc:1197] [/device:CPU:0] (DEBUG INFO) Executor start aborting (this does not indicate an error and you can ignore this message): INVALID_ARGUMENT: You must feed a value for placeholder tensor 'Placeholder/_0' with dtype float and shape [1]
         [[{{node Placeholder/_0}}]]
2023-07-17 23:28:39.807452: W tensorflow/core/grappler/optimizers/data/auto_shard.cc:786] AUTO sharding policy will apply DATA sharding policy as it failed to apply FILE sharding policy because of the following reason: Found an unshardable source dataset: name: "TensorDataset/_2"
op: "TensorDataset"
input: "Placeholder/_0"
input: "Placeholder/_1"
attr {
  key: "Toutput_types"
  value {
    list {
      type: DT_FLOAT
      type: DT_FLOAT
    }
  }
}
attr {
  key: "_cardinality"
  value {
    i: 1
  }
}
attr {
  key: "metadata"
  value {
    s: "\n\017TensorDataset:0"
  }
}
attr {
  key: "output_shapes"
  value {
    list {
      shape {
        dim {
          size: 1
        }
      }
      shape {
        dim {
          size: 1
        }
      }
    }
  }
}
experimental_type {
  type_id: TFT_PRODUCT
  args {
    type_id: TFT_DATASET
    args {
      type_id: TFT_PRODUCT
      args {
        type_id: TFT_TENSOR
        args {
          type_id: TFT_FLOAT
        }
      }
      args {
        type_id: TFT_TENSOR
        args {
          type_id: TFT_FLOAT
        }
      }
    }
  }
}

2023-07-17 23:28:39.823305: I tensorflow/core/common_runtime/executor.cc:1197] [/device:CPU:0] (DEBUG INFO) Executor start aborting (this does not indicate an error and you can ignore this message): INVALID_ARGUMENT: You must feed a value for placeholder tensor 'Placeholder/_1' with dtype float and shape [1]
         [[{{node Placeholder/_1}}]]
2023-07-17 23:28:39.823469: I tensorflow/core/common_runtime/executor.cc:1197] [/device:CPU:0] (DEBUG INFO) Executor start aborting (this does not indicate an error and you can ignore this message): INVALID_ARGUMENT: You must feed a value for placeholder tensor 'Placeholder/_1' with dtype float and shape [1]
         [[{{node Placeholder/_1}}]]
2023-07-17 23:28:40.010671: I tensorflow/core/common_runtime/executor.cc:1197] [/device:CPU:0] (DEBUG INFO) Executor start aborting (this does not indicate an error and you can ignore this message): INVALID_ARGUMENT: You must feed a value for placeholder tensor 'Placeholder/_1' with dtype float and shape [1]
         [[{{node Placeholder/_1}}]]
2023-07-17 23:28:40.010843: I tensorflow/core/common_runtime/executor.cc:1197] [/device:CPU:0] (DEBUG INFO) Executor start aborting (this does not indicate an error and you can ignore this message): INVALID_ARGUMENT: You must feed a value for placeholder tensor 'Placeholder/_0' with dtype float and shape [1]
         [[{{node Placeholder/_0}}]]
2023-07-17 23:28:48.599915: I tensorflow/compiler/xla/service/service.cc:169] XLA service 0x7fe600008ff0 initialized for platform CUDA (this does not guarantee that XLA will be used). Devices:
2023-07-17 23:28:48.599956: I tensorflow/compiler/xla/service/service.cc:177]   StreamExecutor device (0): NVIDIA H100 80GB HBM3, Compute Capability 9.0
2023-07-17 23:28:48.599964: I tensorflow/compiler/xla/service/service.cc:177]   StreamExecutor device (1): NVIDIA H100 80GB HBM3, Compute Capability 9.0
2023-07-17 23:28:48.599971: I tensorflow/compiler/xla/service/service.cc:177]   StreamExecutor device (2): NVIDIA H100 80GB HBM3, Compute Capability 9.0
2023-07-17 23:28:48.599977: I tensorflow/compiler/xla/service/service.cc:177]   StreamExecutor device (3): NVIDIA H100 80GB HBM3, Compute Capability 9.0
2023-07-17 23:28:48.599983: I tensorflow/compiler/xla/service/service.cc:177]   StreamExecutor device (4): NVIDIA H100 80GB HBM3, Compute Capability 9.0
2023-07-17 23:28:48.599988: I tensorflow/compiler/xla/service/service.cc:177]   StreamExecutor device (5): NVIDIA H100 80GB HBM3, Compute Capability 9.0
2023-07-17 23:28:48.599994: I tensorflow/compiler/xla/service/service.cc:177]   StreamExecutor device (6): NVIDIA H100 80GB HBM3, Compute Capability 9.0
2023-07-17 23:28:48.600003: I tensorflow/compiler/xla/service/service.cc:177]   StreamExecutor device (7): NVIDIA H100 80GB HBM3, Compute Capability 9.0
2023-07-17 23:28:48.636269: E tensorflow/compiler/xla/stream_executor/cuda/cuda_dnn.cc:429] Could not create cudnn handle: CUDNN_STATUS_NOT_INITIALIZED
2023-07-17 23:28:48.636335: E tensorflow/compiler/xla/stream_executor/cuda/cuda_dnn.cc:438] Possibly insufficient driver version: 535.54.3
2023-07-17 23:28:48.649063: E tensorflow/compiler/xla/stream_executor/cuda/cuda_dnn.cc:429] Could not create cudnn handle: CUDNN_STATUS_NOT_INITIALIZED
2023-07-17 23:28:48.649131: E tensorflow/compiler/xla/stream_executor/cuda/cuda_dnn.cc:438] Possibly insufficient driver version: 535.54.3
2023-07-17 23:28:48.711026: I tensorflow/compiler/xla/stream_executor/cuda/cuda_dnn.cc:424] Loaded cuDNN version 8600
2023-07-17 23:28:48.712282: E tensorflow/compiler/xla/status_macros.cc:57] INTERNAL: RET_CHECK failure (tensorflow/compiler/xla/service/gpu/gpu_compiler.cc:618) dnn != nullptr
*** Begin stack trace ***
        tsl::CurrentStackTrace[abi:cxx11]()

        xla::status_macros::MakeErrorStream::Impl::GetStatus()
        xla::gpu::GpuCompiler::OptimizeHloModule(xla::HloModule*, stream_executor::StreamExecutor*, stream_executor::DeviceMemoryAllocator*, xla::gpu::GpuTargetConfig const&, xla::AutotuneResults const*)
        xla::gpu::GpuCompiler::RunHloPasses(std::unique_ptr<xla::HloModule, std::default_delete<xla::HloModule> >, stream_executor::StreamExecutor*, xla::Compiler::CompileOptions const&)
        xla::Service::BuildExecutable(xla::HloModuleProto const&, std::unique_ptr<xla::HloModuleConfig, std::default_delete<xla::HloModuleConfig> >, xla::Backend*, stream_executor::StreamExecutor*, xla::Compiler::CompileOptions const&, bool)
        xla::LocalService::CompileExecutables(xla::XlaComputation const&, absl::lts_20220623::Span<xla::Shape const* const>, xla::ExecutableBuildOptions const&)
        xla::LocalClient::Compile(xla::XlaComputation const&, absl::lts_20220623::Span<xla::Shape const* const>, xla::ExecutableBuildOptions const&)
        tensorflow::XlaDeviceCompilerClient::BuildExecutable(tensorflow::XlaCompiler::Options const&, tensorflow::XlaCompilationResult const&)
        tensorflow::DeviceCompiler<xla::LocalExecutable, xla::LocalClient>::CompileStrict(tensorflow::DeviceCompilationClusterSignature const&, tensorflow::XlaCompiler::CompileOptions const&, tensorflow::XlaCompiler::Options const&, std::vector<tensorflow::XlaArgument, std::allocator<tensorflow::XlaArgument> > const&, tensorflow::NameAttrList const&, tensorflow::DeviceCompilationCache<xla::LocalExecutable>::Value, tensorflow::DeviceCompiler<xla::LocalExecutable, xla::LocalClient>::CompileScope, tensorflow::OpKernelContext*, tensorflow::DeviceCompilationProfiler*, tsl::mutex*)
        tensorflow::DeviceCompiler<xla::LocalExecutable, xla::LocalClient>::CompileImpl(tensorflow::XlaCompiler::CompileOptions const&, tensorflow::XlaCompiler::Options const&, tensorflow::NameAttrList const&, std::vector<tensorflow::XlaArgument, std::allocator<tensorflow::XlaArgument> > const&, tensorflow::DeviceCompiler<xla::LocalExecutable, xla::LocalClient>::CompileScope, tensorflow::DeviceCompileMode, tensorflow::OpKernelContext*, tensorflow::DeviceCompilationProfiler*, tensorflow::XlaCompilationResult const**, xla::LocalExecutable**)

        tensorflow::XlaLocalLaunchBase::ComputeAsync(tensorflow::OpKernelContext*, std::function<void ()>)
        tensorflow::BaseGPUDevice::ComputeAsync(tensorflow::AsyncOpKernel*, tensorflow::OpKernelContext*, std::function<void ()>)


        Eigen::ThreadPoolTempl<tsl::thread::EigenEnvironment>::WorkerLoop(int)
        std::_Function_handler<void (), tsl::thread::EigenEnvironment::CreateThread(std::function<void ()>)::{lambda()#1}>::_M_invoke(std::_Any_data const&)



*** End stack trace ***

2023-07-17 23:28:48.712483: W tensorflow/core/framework/op_kernel.cc:1830] OP_REQUIRES failed at xla_ops.cc:362 : INTERNAL: RET_CHECK failure (tensorflow/compiler/xla/service/gpu/gpu_compiler.cc:618) dnn != nullptr
2023-07-17 23:28:48.712506: I tensorflow/core/common_runtime/executor.cc:1197] [/job:localhost/replica:0/task:0/device:GPU:0] (DEBUG INFO) Executor start aborting (this does not indicate an error and you can ignore this message): INTERNAL: RET_CHECK failure (tensorflow/compiler/xla/service/gpu/gpu_compiler.cc:618) dnn != nullptr
         [[{{node update_0_1/StatefulPartitionedCall}}]]
2023-07-17 23:28:48.727426: E tensorflow/compiler/xla/status_macros.cc:57] INTERNAL: RET_CHECK failure (tensorflow/compiler/xla/service/gpu/gpu_compiler.cc:618) dnn != nullptr
*** Begin stack trace ***
        tsl::CurrentStackTrace[abi:cxx11]()

        xla::status_macros::MakeErrorStream::Impl::GetStatus()
        xla::gpu::GpuCompiler::OptimizeHloModule(xla::HloModule*, stream_executor::StreamExecutor*, stream_executor::DeviceMemoryAllocator*, xla::gpu::GpuTargetConfig const&, xla::AutotuneResults const*)
        xla::gpu::GpuCompiler::RunHloPasses(std::unique_ptr<xla::HloModule, std::default_delete<xla::HloModule> >, stream_executor::StreamExecutor*, xla::Compiler::CompileOptions const&)
        xla::Service::BuildExecutable(xla::HloModuleProto const&, std::unique_ptr<xla::HloModuleConfig, std::default_delete<xla::HloModuleConfig> >, xla::Backend*, stream_executor::StreamExecutor*, xla::Compiler::CompileOptions const&, bool)
        xla::LocalService::CompileExecutables(xla::XlaComputation const&, absl::lts_20220623::Span<xla::Shape const* const>, xla::ExecutableBuildOptions const&)
        xla::LocalClient::Compile(xla::XlaComputation const&, absl::lts_20220623::Span<xla::Shape const* const>, xla::ExecutableBuildOptions const&)
        tensorflow::XlaDeviceCompilerClient::BuildExecutable(tensorflow::XlaCompiler::Options const&, tensorflow::XlaCompilationResult const&)
        tensorflow::DeviceCompiler<xla::LocalExecutable, xla::LocalClient>::CompileStrict(tensorflow::DeviceCompilationClusterSignature const&, tensorflow::XlaCompiler::CompileOptions const&, tensorflow::XlaCompiler::Options const&, std::vector<tensorflow::XlaArgument, std::allocator<tensorflow::XlaArgument> > const&, tensorflow::NameAttrList const&, tensorflow::DeviceCompilationCache<xla::LocalExecutable>::Value, tensorflow::DeviceCompiler<xla::LocalExecutable, xla::LocalClient>::CompileScope, tensorflow::OpKernelContext*, tensorflow::DeviceCompilationProfiler*, tsl::mutex*)
        tensorflow::DeviceCompiler<xla::LocalExecutable, xla::LocalClient>::CompileImpl(tensorflow::XlaCompiler::CompileOptions const&, tensorflow::XlaCompiler::Options const&, tensorflow::NameAttrList const&, std::vector<tensorflow::XlaArgument, std::allocator<tensorflow::XlaArgument> > const&, tensorflow::DeviceCompiler<xla::LocalExecutable, xla::LocalClient>::CompileScope, tensorflow::DeviceCompileMode, tensorflow::OpKernelContext*, tensorflow::DeviceCompilationProfiler*, tensorflow::XlaCompilationResult const**, xla::LocalExecutable**)

        tensorflow::XlaLocalLaunchBase::ComputeAsync(tensorflow::OpKernelContext*, std::function<void ()>)
        tensorflow::BaseGPUDevice::ComputeAsync(tensorflow::AsyncOpKernel*, tensorflow::OpKernelContext*, std::function<void ()>)


        Eigen::ThreadPoolTempl<tsl::thread::EigenEnvironment>::WorkerLoop(int)
        std::_Function_handler<void (), tsl::thread::EigenEnvironment::CreateThread(std::function<void ()>)::{lambda()#1}>::_M_invoke(std::_Any_data const&)



*** End stack trace ***

2023-07-17 23:28:48.727590: W tensorflow/core/framework/op_kernel.cc:1830] OP_REQUIRES failed at xla_ops.cc:362 : INTERNAL: RET_CHECK failure (tensorflow/compiler/xla/service/gpu/gpu_compiler.cc:618) dnn != nullptr
2023-07-17 23:28:48.727609: I tensorflow/core/common_runtime/executor.cc:1197] [/job:localhost/replica:0/task:0/device:GPU:1] (DEBUG INFO) Executor start aborting (this does not indicate an error and you can ignore this message): INTERNAL: RET_CHECK failure (tensorflow/compiler/xla/service/gpu/gpu_compiler.cc:618) dnn != nullptr
         [[{{node update_1_1/StatefulPartitionedCall}}]]
2023-07-17 23:28:48.728128: E tensorflow/compiler/xla/stream_executor/cuda/cuda_dnn.cc:429] Could not create cudnn handle: CUDNN_STATUS_NOT_INITIALIZED
2023-07-17 23:28:48.728185: E tensorflow/compiler/xla/stream_executor/cuda/cuda_dnn.cc:438] Possibly insufficient driver version: 535.54.3
2023-07-17 23:28:48.740946: W tensorflow/compiler/xla/stream_executor/gpu/asm_compiler.cc:231] Falling back to the CUDA driver for PTX compilation; ptxas does not support CC 9.0
2023-07-17 23:28:48.740965: W tensorflow/compiler/xla/stream_executor/gpu/asm_compiler.cc:234] Used ptxas at ptxas
2023-07-17 23:28:48.742906: E tensorflow/compiler/xla/stream_executor/cuda/cuda_dnn.cc:429] Could not create cudnn handle: CUDNN_STATUS_NOT_INITIALIZED
2023-07-17 23:28:48.742961: E tensorflow/compiler/xla/stream_executor/cuda/cuda_dnn.cc:438] Possibly insufficient driver version: 535.54.3
2023-07-17 23:28:48.770994: I ./tensorflow/compiler/jit/device_compiler.h:180] Compiled cluster using XLA!  This line is logged at most once for the lifetime of the process.
2023-07-17 23:28:48.806579: E tensorflow/compiler/xla/status_macros.cc:57] INTERNAL: RET_CHECK failure (tensorflow/compiler/xla/service/gpu/gpu_compiler.cc:618) dnn != nullptr
*** Begin stack trace ***
        tsl::CurrentStackTrace[abi:cxx11]()

        xla::status_macros::MakeErrorStream::Impl::GetStatus()
        xla::gpu::GpuCompiler::OptimizeHloModule(xla::HloModule*, stream_executor::StreamExecutor*, stream_executor::DeviceMemoryAllocator*, xla::gpu::GpuTargetConfig const&, xla::AutotuneResults const*)
        xla::gpu::GpuCompiler::RunHloPasses(std::unique_ptr<xla::HloModule, std::default_delete<xla::HloModule> >, stream_executor::StreamExecutor*, xla::Compiler::CompileOptions const&)
        xla::Service::BuildExecutable(xla::HloModuleProto const&, std::unique_ptr<xla::HloModuleConfig, std::default_delete<xla::HloModuleConfig> >, xla::Backend*, stream_executor::StreamExecutor*, xla::Compiler::CompileOptions const&, bool)
        xla::LocalService::CompileExecutables(xla::XlaComputation const&, absl::lts_20220623::Span<xla::Shape const* const>, xla::ExecutableBuildOptions const&)
        xla::LocalClient::Compile(xla::XlaComputation const&, absl::lts_20220623::Span<xla::Shape const* const>, xla::ExecutableBuildOptions const&)
        tensorflow::XlaDeviceCompilerClient::BuildExecutable(tensorflow::XlaCompiler::Options const&, tensorflow::XlaCompilationResult const&)
        tensorflow::DeviceCompiler<xla::LocalExecutable, xla::LocalClient>::CompileStrict(tensorflow::DeviceCompilationClusterSignature const&, tensorflow::XlaCompiler::CompileOptions const&, tensorflow::XlaCompiler::Options const&, std::vector<tensorflow::XlaArgument, std::allocator<tensorflow::XlaArgument> > const&, tensorflow::NameAttrList const&, tensorflow::DeviceCompilationCache<xla::LocalExecutable>::Value, tensorflow::DeviceCompiler<xla::LocalExecutable, xla::LocalClient>::CompileScope, tensorflow::OpKernelContext*, tensorflow::DeviceCompilationProfiler*, tsl::mutex*)
        tensorflow::DeviceCompiler<xla::LocalExecutable, xla::LocalClient>::CompileImpl(tensorflow::XlaCompiler::CompileOptions const&, tensorflow::XlaCompiler::Options const&, tensorflow::NameAttrList const&, std::vector<tensorflow::XlaArgument, std::allocator<tensorflow::XlaArgument> > const&, tensorflow::DeviceCompiler<xla::LocalExecutable, xla::LocalClient>::CompileScope, tensorflow::DeviceCompileMode, tensorflow::OpKernelContext*, tensorflow::DeviceCompilationProfiler*, tensorflow::XlaCompilationResult const**, xla::LocalExecutable**)

        tensorflow::XlaLocalLaunchBase::ComputeAsync(tensorflow::OpKernelContext*, std::function<void ()>)
        tensorflow::BaseGPUDevice::ComputeAsync(tensorflow::AsyncOpKernel*, tensorflow::OpKernelContext*, std::function<void ()>)


        Eigen::ThreadPoolTempl<tsl::thread::EigenEnvironment>::WorkerLoop(int)
        std::_Function_handler<void (), tsl::thread::EigenEnvironment::CreateThread(std::function<void ()>)::{lambda()#1}>::_M_invoke(std::_Any_data const&)



*** End stack trace ***

2023-07-17 23:28:48.806809: W tensorflow/core/framework/op_kernel.cc:1830] OP_REQUIRES failed at xla_ops.cc:362 : INTERNAL: RET_CHECK failure (tensorflow/compiler/xla/service/gpu/gpu_compiler.cc:618) dnn != nullptr
2023-07-17 23:28:48.818268: E tensorflow/compiler/xla/status_macros.cc:57] INTERNAL: RET_CHECK failure (tensorflow/compiler/xla/service/gpu/gpu_compiler.cc:618) dnn != nullptr
*** Begin stack trace ***
        tsl::CurrentStackTrace[abi:cxx11]()

        xla::status_macros::MakeErrorStream::Impl::GetStatus()
        xla::gpu::GpuCompiler::OptimizeHloModule(xla::HloModule*, stream_executor::StreamExecutor*, stream_executor::DeviceMemoryAllocator*, xla::gpu::GpuTargetConfig const&, xla::AutotuneResults const*)
        xla::gpu::GpuCompiler::RunHloPasses(std::unique_ptr<xla::HloModule, std::default_delete<xla::HloModule> >, stream_executor::StreamExecutor*, xla::Compiler::CompileOptions const&)
        xla::Service::BuildExecutable(xla::HloModuleProto const&, std::unique_ptr<xla::HloModuleConfig, std::default_delete<xla::HloModuleConfig> >, xla::Backend*, stream_executor::StreamExecutor*, xla::Compiler::CompileOptions const&, bool)
        xla::LocalService::CompileExecutables(xla::XlaComputation const&, absl::lts_20220623::Span<xla::Shape const* const>, xla::ExecutableBuildOptions const&)
        xla::LocalClient::Compile(xla::XlaComputation const&, absl::lts_20220623::Span<xla::Shape const* const>, xla::ExecutableBuildOptions const&)
        tensorflow::XlaDeviceCompilerClient::BuildExecutable(tensorflow::XlaCompiler::Options const&, tensorflow::XlaCompilationResult const&)
        tensorflow::DeviceCompiler<xla::LocalExecutable, xla::LocalClient>::CompileStrict(tensorflow::DeviceCompilationClusterSignature const&, tensorflow::XlaCompiler::CompileOptions const&, tensorflow::XlaCompiler::Options const&, std::vector<tensorflow::XlaArgument, std::allocator<tensorflow::XlaArgument> > const&, tensorflow::NameAttrList const&, tensorflow::DeviceCompilationCache<xla::LocalExecutable>::Value, tensorflow::DeviceCompiler<xla::LocalExecutable, xla::LocalClient>::CompileScope, tensorflow::OpKernelContext*, tensorflow::DeviceCompilationProfiler*, tsl::mutex*)
        tensorflow::DeviceCompiler<xla::LocalExecutable, xla::LocalClient>::CompileImpl(tensorflow::XlaCompiler::CompileOptions const&, tensorflow::XlaCompiler::Options const&, tensorflow::NameAttrList const&, std::vector<tensorflow::XlaArgument, std::allocator<tensorflow::XlaArgument> > const&, tensorflow::DeviceCompiler<xla::LocalExecutable, xla::LocalClient>::CompileScope, tensorflow::DeviceCompileMode, tensorflow::OpKernelContext*, tensorflow::DeviceCompilationProfiler*, tensorflow::XlaCompilationResult const**, xla::LocalExecutable**)

        tensorflow::XlaLocalLaunchBase::ComputeAsync(tensorflow::OpKernelContext*, std::function<void ()>)
        tensorflow::BaseGPUDevice::ComputeAsync(tensorflow::AsyncOpKernel*, tensorflow::OpKernelContext*, std::function<void ()>)


        Eigen::ThreadPoolTempl<tsl::thread::EigenEnvironment>::WorkerLoop(int)
        std::_Function_handler<void (), tsl::thread::EigenEnvironment::CreateThread(std::function<void ()>)::{lambda()#1}>::_M_invoke(std::_Any_data const&)



*** End stack trace ***

2023-07-17 23:28:48.818437: W tensorflow/core/framework/op_kernel.cc:1830] OP_REQUIRES failed at xla_ops.cc:362 : INTERNAL: RET_CHECK failure (tensorflow/compiler/xla/service/gpu/gpu_compiler.cc:618) dnn != nullptr
2023-07-17 23:28:48.858099: I tensorflow/core/common_runtime/executor.cc:1197] [/job:localhost/replica:0/task:0/device:GPU:2] (DEBUG INFO) Executor start aborting (this does not indicate an error and you can ignore this message): INTERNAL: RET_CHECK failure (tensorflow/compiler/xla/service/gpu/gpu_compiler.cc:618) dnn != nullptr
         [[{{node update_0_1/StatefulPartitionedCall}}]]
         [[GroupCrossDeviceControlEdges_1/Identity_7/_166]]
Traceback (most recent call last):
  File "/home/user4/project/src/test.py", line 8, in <module>
    model.fit(dataset)
  File "/home/user4/miniconda3/envs/tf/lib/python3.9/site-packages/keras/utils/traceback_utils.py", line 70, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/home/user4/miniconda3/envs/tf/lib/python3.9/site-packages/tensorflow/python/eager/execute.py", line 52, in quick_execute
    tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
tensorflow.python.framework.errors_impl.InternalError: Graph execution error:

Detected at node 'update_1_1/StatefulPartitionedCall' defined at (most recent call last):
    File "/home/user4/project/src/test.py", line 8, in <module>
      model.fit(dataset)
    File "/home/user4/miniconda3/envs/tf/lib/python3.9/site-packages/keras/utils/traceback_utils.py", line 65, in error_handler
      return fn(*args, **kwargs)
    File "/home/user4/miniconda3/envs/tf/lib/python3.9/site-packages/keras/engine/training.py", line 1685, in fit
      tmp_logs = self.train_function(iterator)
    File "/home/user4/miniconda3/envs/tf/lib/python3.9/site-packages/keras/engine/training.py", line 1284, in train_function
      return step_function(self, iterator)
    File "/home/user4/miniconda3/envs/tf/lib/python3.9/site-packages/keras/engine/training.py", line 1268, in step_function
      outputs = model.distribute_strategy.run(run_step, args=(data,))
    File "/home/user4/miniconda3/envs/tf/lib/python3.9/site-packages/keras/optimizers/optimizer.py", line 1250, in _distributed_apply_gradients_fn
      distribution.extended.update(
    File "/home/user4/miniconda3/envs/tf/lib/python3.9/site-packages/keras/optimizers/optimizer.py", line 1245, in apply_grad_to_update_var
      return self._update_step_xla(grad, var, id(self._var_key(var)))
Node: 'update_1_1/StatefulPartitionedCall'
Detected at node 'update_0_1/StatefulPartitionedCall' defined at (most recent call last):
    File "/home/user4/project/src/test.py", line 8, in <module>
      model.fit(dataset)
    File "/home/user4/miniconda3/envs/tf/lib/python3.9/site-packages/keras/utils/traceback_utils.py", line 65, in error_handler
      return fn(*args, **kwargs)
    File "/home/user4/miniconda3/envs/tf/lib/python3.9/site-packages/keras/engine/training.py", line 1685, in fit
      tmp_logs = self.train_function(iterator)
    File "/home/user4/miniconda3/envs/tf/lib/python3.9/site-packages/keras/engine/training.py", line 1284, in train_function
      return step_function(self, iterator)
    File "/home/user4/miniconda3/envs/tf/lib/python3.9/site-packages/keras/engine/training.py", line 1268, in step_function
      outputs = model.distribute_strategy.run(run_step, args=(data,))
    File "/home/user4/miniconda3/envs/tf/lib/python3.9/site-packages/keras/optimizers/optimizer.py", line 1250, in _distributed_apply_gradients_fn
      distribution.extended.update(
    File "/home/user4/miniconda3/envs/tf/lib/python3.9/site-packages/keras/optimizers/optimizer.py", line 1245, in apply_grad_to_update_var
      return self._update_step_xla(grad, var, id(self._var_key(var)))
Node: 'update_0_1/StatefulPartitionedCall'
Detected at node 'update_0_1/StatefulPartitionedCall' defined at (most recent call last):
    File "/home/user4/project/src/test.py", line 8, in <module>
      model.fit(dataset)
    File "/home/user4/miniconda3/envs/tf/lib/python3.9/site-packages/keras/utils/traceback_utils.py", line 65, in error_handler
      return fn(*args, **kwargs)
    File "/home/user4/miniconda3/envs/tf/lib/python3.9/site-packages/keras/engine/training.py", line 1685, in fit
      tmp_logs = self.train_function(iterator)
    File "/home/user4/miniconda3/envs/tf/lib/python3.9/site-packages/keras/engine/training.py", line 1284, in train_function
      return step_function(self, iterator)
    File "/home/user4/miniconda3/envs/tf/lib/python3.9/site-packages/keras/engine/training.py", line 1268, in step_function
      outputs = model.distribute_strategy.run(run_step, args=(data,))
    File "/home/user4/miniconda3/envs/tf/lib/python3.9/site-packages/keras/optimizers/optimizer.py", line 1250, in _distributed_apply_gradients_fn
      distribution.extended.update(
    File "/home/user4/miniconda3/envs/tf/lib/python3.9/site-packages/keras/optimizers/optimizer.py", line 1245, in apply_grad_to_update_var
      return self._update_step_xla(grad, var, id(self._var_key(var)))
Node: 'update_0_1/StatefulPartitionedCall'
3 root error(s) found.
  (0) INTERNAL:  RET_CHECK failure (tensorflow/compiler/xla/service/gpu/gpu_compiler.cc:618) dnn != nullptr
         [[{{node update_1_1/StatefulPartitionedCall}}]]
  (1) INTERNAL:  RET_CHECK failure (tensorflow/compiler/xla/service/gpu/gpu_compiler.cc:618) dnn != nullptr
         [[{{node update_0_1/StatefulPartitionedCall}}]]
         [[GroupCrossDeviceControlEdges_1/Identity_7/_166]]
  (2) INTERNAL:  RET_CHECK failure (tensorflow/compiler/xla/service/gpu/gpu_compiler.cc:618) dnn != nullptr
         [[{{node update_0_1/StatefulPartitionedCall}}]]
0 successful operations.
0 derived errors ignored. [Op:__inference_train_function_1810]
```
