CUDA_ERROR_ILLEGAL_ADDRESS
In this gist, I have removed the dependency on `transformers` and a large amount of irrelevant code:
V1: https://gist.github.com/AndreasMadsen/dc5785b5a55bf740c555b2fb5cdab1db (~600 LOC)
V2: https://gist.github.com/AndreasMadsen/5fdaa8431929e25cf3a990f234f88a8c (update, more code removed. ~450 LOC)
V3: https://gist.github.com/AndreasMadsen/2590423c055bc47a932a449a5161bac7 (update, more code removed.  209 LOC)

I have tried to reduce the code further, but this causes the bug to disappear. @pjpratik I think you should add the `comp:xla` label too. As this bug only occurs when XLA is enabled. There is no failure with the reproduction script on `2.12.0rc0`. However, given how difficult this bug is to reproduce (e.g. changing just the seed of the shuffle function will make it disappear). I can't tell if this solves the issue consistently. The error appears to be specific to V100, as it does not reproduce on T4, A100, and P100 GPUs. I have however reproduced it on two different systems running V100. It is possible this is related to the specific memory size of the GPU. Hi, @AndreasMadsen 

Apologize for the delay and Thank you for your detailed observations and efforts to find out the root cause for this issue and We'll try it from our end and will check whether this issue is specific to V100 GPU or not and we'll update you soon here. Thank you for noticing this issue. Thank you! @gaikwadrahul8 Thanks for looking into this. It did take me several days to make this error reproducible. So I hope it will help you. I suspect it's a rather complex integration bug, as removing pretty much any line in the latest (V3) script cause the bug to disappear.  Hi, @AndreasMadsen 

Apologize for the delay and I was able to replicate the issue on `Tesla V100 GPU` and I was getting the exact same error message which you mentioned in your error log output above with `tensorflow/tensorflow:latest-gpu` docker image which is using `Tensorflow==2.11`, I have added screenshot for your reference below and even tested the same code on Google Colab and it's working fine as expected here is [gist-file](https://colab.research.google.com/gist/gaikwadrahul8/caf4e40e4ef5c295f9c0b1defcacc63a/-59671.ipynb), I did not run for complete epoch because it was taking more time but code was running as expected but as you confirmed it's working fine on all other GPU's so this issue seems like specific to `Tesla V100 GPU` at the moment. Thank you!

When I tried the same code with `tensorflow/tensorflow:2.12.0rc0-gpu` and `tensorflow/tensorflow:nightly-gpu`  on `Tesla V100 GPU` it's working as expected, I have added screenshot below

Here is screenshot for code execution with `tensorflow/tensorflow:latest-gpu` docker image

![image](https://user-images.githubusercontent.com/115997457/220920902-6a6004fb-297f-4bbb-8683-2a6e14d1b83a.png)

Here is screenshot for code execution on Google-Colab :

![image](https://user-images.githubusercontent.com/115997457/220929005-86bab24b-ef55-4e61-90fe-8cf95fb8a5cf.png)

Here is screenshot for code execution with `tensorflow/tensorflow:2.12.0rc0-gpu` docker image:

![image](https://user-images.githubusercontent.com/115997457/220937369-56b79ed5-d6db-4f28-9082-19d52762c4ad.png)

Here is screenshot for code execution with `tensorflow/tensorflow:nightly-gpu` docker image:

![image](https://user-images.githubusercontent.com/115997457/220945432-a6f48ee1-afa3-426b-b793-d1f0db823e37.png)


 Hi @gaikwadrahul8, you didn't ask a question, so I'm not sure what you want me to respond to.

I agree with your findings. Given how sensitive the bug is in terms of reproduction, I would not trust that it is something specific to V100 TF 2.11. I hope you will be able to debug this further and get to the resolve bug. Note, that this kind of bug has been reported quite a few times before, but I think most issues were not able to provide a reproduction script. So I hope you will find this script quite useful.

ISSUE: https://github.com/tensorflow/tensorflow/issues/50735

I remember there were more closed issues, when I was researching solutions.

 Hi. @AndreasMadsen 

Thank you for mentioning the similar issue, It seems like we have to dig more into this issue to find out the root cause and we'll update you soon and I was able to replicate the exact same issue with error log with `tensorflow/tensorflow:latest-gpu ` docker image that's the reason I did not ask any questions to you if we need any input or further details from your end certainly we will ask you. Thank you for noticing this issue, I really appreciate your efforts and time. Thank you! @gaikwadrahul8 Thanks for letting me know. I'm was just confused about the `stat:awaiting response` label. Hi, @sachinprasadhs 

Could you please look into this issue ? Thank you! Hi @sachinprasadhs, did you manage to investigate this?  Hi @AndreasMadsen , Apologies for the delayed response, based on the issue thread so far, it is difficult to point to an area why this is happening.
But you could try below steps and see if that works.
1. Try reinstalling the CUDA and cuDNN dependency, in some cases it could be due to improper installation of these dependencies.
2. Try setting the environmental variable `TF_FORCE_GPU_ALLOW_GROWTH `to `true`

If the above steps did not work could you please run the model with the env var `CUDA_LAUNCH_BLOCKING` set to `1` and attach the full logs? @sachinprasadhs Thanks for your response. Please note that I spent several days reducing this to a script that reproduces the issue. I think you can benifit tremendously from using that. Furthermore, @gaikwadrahul8 confirms that the issue could be reproduced with a V100 GPU.

1. Try reinstalling the CUDA and cuDNN dependency, in some cases it could be due to improper installation of these dependencies.

Let's not go down that route, since @gaikwadrahul8 was able to reproduce the issue on Google's infrastructure.

2. Try setting the environmental variable `TF_FORCE_GPU_ALLOW_GROWTH` to `true`.

Same issue. Expand the `<details>` to see the log.

<details>
<summary><code>TF_FORCE_GPU_ALLOW_GROWTH=true python main.py</code></summary>

```
2023-03-15 11:55:41.321502: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 AVX512F AVX512_VNNI FMA
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
2023-03-15 11:55:41.504643: I tensorflow/core/util/port.cc:104] 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-03-15 11:55:42.479709: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer.so.7'; dlerror: libnvinfer.so.7: cannot open shared object file: No such file or directory
2023-03-15 11:55:42.480077: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer_plugin.so.7'; dlerror: libnvinfer_plugin.so.7: cannot open shared object file: No such file or directory
2023-03-15 11:55:42.480113: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.
2023-03-15 11:55:43.729040: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 AVX512F AVX512_VNNI FMA
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
2023-03-15 11:55:44.593353: W tensorflow/core/common_runtime/gpu/gpu_bfc_allocator.cc:42] Overriding orig_value setting because the TF_FORCE_GPU_ALLOW_GROWTH environment variable is set. Original config value was 0.
2023-03-15 11:55:44.593541: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1613] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 30971 MB memory:  -> device: 0, name: Tesla V100-SXM2-32GB, pci bus id: 0000:18:00.0, compute capability: 7.0
2023-03-15 11:56:14.842711: I tensorflow/compiler/xla/service/service.cc:173] XLA service 0x20192630 initialized for platform CUDA (this does not guarantee that XLA will be used). Devices:
2023-03-15 11:56:14.842838: I tensorflow/compiler/xla/service/service.cc:181]   StreamExecutor device (0): Tesla V100-SXM2-32GB, Compute Capability 7.0
2023-03-15 11:56:15.808448: I tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.cc:268] disabling MLIR crash reproducer, set env var `MLIR_CRASH_REPRODUCER_DIRECTORY` to enable.
2023-03-15 11:56:15.816882: W tensorflow/compiler/tf2xla/kernels/random_ops.cc:57] Warning: Using tf.random.uniform with XLA compilation will ignore seeds; consider using tf.random.stateless_uniform instead if reproducible behavior is desired. tf_roberta_for_sequence_classification/layer_._0/self/dropout/dropout/random_uniform/RandomUniform
2023-03-15 11:57:26.919108: I tensorflow/compiler/xla/stream_executor/gpu/asm_compiler.cc:325] ptxas warning : Registers are spilled to local memory in function 'sort_9'
ptxas warning : Registers are spilled to local memory in function 'fusion_1107'

2023-03-15 11:57:27.290497: I tensorflow/compiler/jit/xla_compilation_cache.cc:477] Compiled cluster using XLA!  This line is logged at most once for the lifetime of the process.
2023-03-15 11:57:27.774582: E tensorflow/compiler/xla/stream_executor/cuda/cuda_driver.cc:1159] failed to enqueue async memcpy from device to host: CUDA_ERROR_ILLEGAL_ADDRESS: an illegal memory access was encountered; host dst: 0x2b4b7e184eeb; GPU src: 0x2b4e0836e500; size: 1=0x1
2023-03-15 11:57:27.774751: I tensorflow/compiler/xla/stream_executor/stream.cc:2485] INTERNAL: Unknown error
2023-03-15 11:57:27.774869: I tensorflow/compiler/xla/stream_executor/stream.cc:2489] [stream=0x1c1751e0,impl=0x5e7c020] INTERNAL: stream did not block host until done; was already in an error state
Traceback (most recent call last):
  File "/scratch/anmadc/workspace/economical-roar/main.py", line 209, in <module>
    model.fit(dataset_train_batched, verbose=1, epochs=1)
  File "/localscratch/anmadc.62577500.0/env/lib/python3.10/site-packages/keras/utils/traceback_utils.py", line 70, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/localscratch/anmadc.62577500.0/env/lib/python3.10/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 'StatefulPartitionedCall' defined at (most recent call last):
    File "/scratch/anmadc/workspace/economical-roar/main.py", line 209, in <module>
      model.fit(dataset_train_batched, verbose=1, epochs=1)
    File "/localscratch/anmadc.62577500.0/env/lib/python3.10/site-packages/keras/utils/traceback_utils.py", line 65, in error_handler
      return fn(*args, **kwargs)
    File "/localscratch/anmadc.62577500.0/env/lib/python3.10/site-packages/keras/engine/training.py", line 1650, in fit
      tmp_logs = self.train_function(iterator)
    File "/localscratch/anmadc.62577500.0/env/lib/python3.10/site-packages/keras/engine/training.py", line 1249, in train_function
      return step_function(self, iterator)
    File "/localscratch/anmadc.62577500.0/env/lib/python3.10/site-packages/keras/engine/training.py", line 1233, in step_function
      outputs = model.distribute_strategy.run(run_step, args=(data,))
Node: 'StatefulPartitionedCall'
Failed to retrieve branch_index value on stream 0x1c1751e0: stream did not block host until done; was already in an error state.
	 [[{{node StatefulPartitionedCall}}]] [Op:__inference_train_function_37296]
```
</details>


<details>
<summary><code>CUDA_LAUNCH_BLOCKING=1 TF_FORCE_GPU_ALLOW_GROWTH=true python main.py</code></summary>

```
2023-03-15 11:57:59.121539: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 AVX512F AVX512_VNNI FMA
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
2023-03-15 11:57:59.307027: I tensorflow/core/util/port.cc:104] 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-03-15 11:58:00.281216: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer.so.7'; dlerror: libnvinfer.so.7: cannot open shared object file: No such file or directory
2023-03-15 11:58:00.281601: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer_plugin.so.7'; dlerror: libnvinfer_plugin.so.7: cannot open shared object file: No such file or directory
2023-03-15 11:58:00.281646: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.
2023-03-15 11:58:01.543891: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 AVX512F AVX512_VNNI FMA
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
2023-03-15 11:58:02.412447: W tensorflow/core/common_runtime/gpu/gpu_bfc_allocator.cc:42] Overriding orig_value setting because the TF_FORCE_GPU_ALLOW_GROWTH environment variable is set. Original config value was 0.
2023-03-15 11:58:02.412647: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1613] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 30971 MB memory:  -> device: 0, name: Tesla V100-SXM2-32GB, pci bus id: 0000:18:00.0, compute capability: 7.0
2023-03-15 11:58:32.523032: I tensorflow/compiler/xla/service/service.cc:173] XLA service 0x1fb485b0 initialized for platform CUDA (this does not guarantee that XLA will be used). Devices:
2023-03-15 11:58:32.523516: I tensorflow/compiler/xla/service/service.cc:181]   StreamExecutor device (0): Tesla V100-SXM2-32GB, Compute Capability 7.0
2023-03-15 11:58:33.492464: I tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.cc:268] disabling MLIR crash reproducer, set env var `MLIR_CRASH_REPRODUCER_DIRECTORY` to enable.
2023-03-15 11:58:33.500937: W tensorflow/compiler/tf2xla/kernels/random_ops.cc:57] Warning: Using tf.random.uniform with XLA compilation will ignore seeds; consider using tf.random.stateless_uniform instead if reproducible behavior is desired. tf_roberta_for_sequence_classification/layer_._0/self/dropout/dropout/random_uniform/RandomUniform
2023-03-15 11:59:44.579959: I tensorflow/compiler/xla/stream_executor/gpu/asm_compiler.cc:325] ptxas warning : Registers are spilled to local memory in function 'sort_9'
ptxas warning : Registers are spilled to local memory in function 'fusion_1107'

2023-03-15 11:59:44.945238: I tensorflow/compiler/jit/xla_compilation_cache.cc:477] Compiled cluster using XLA!  This line is logged at most once for the lifetime of the process.
Traceback (most recent call last):
  File "/scratch/anmadc/workspace/economical-roar/main.py", line 209, in <module>
    model.fit(dataset_train_batched, verbose=1, epochs=1)
  File "/localscratch/anmadc.62577500.0/env/lib/python3.10/site-packages/keras/utils/traceback_utils.py", line 70, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/localscratch/anmadc.62577500.0/env/lib/python3.10/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 'StatefulPartitionedCall' defined at (most recent call last):
    File "/scratch/anmadc/workspace/economical-roar/main.py", line 209, in <module>
      model.fit(dataset_train_batched, verbose=1, epochs=1)
    File "/localscratch/anmadc.62577500.0/env/lib/python3.10/site-packages/keras/utils/traceback_utils.py", line 65, in error_handler
      return fn(*args, **kwargs)
    File "/localscratch/anmadc.62577500.0/env/lib/python3.10/site-packages/keras/engine/training.py", line 1650, in fit
      tmp_logs = self.train_function(iterator)
    File "/localscratch/anmadc.62577500.0/env/lib/python3.10/site-packages/keras/engine/training.py", line 1249, in train_function
      return step_function(self, iterator)
    File "/localscratch/anmadc.62577500.0/env/lib/python3.10/site-packages/keras/engine/training.py", line 1233, in step_function
      outputs = model.distribute_strategy.run(run_step, args=(data,))
Node: 'StatefulPartitionedCall'
Failed to launch CUDA kernel: fusion_1428 with block dimensions: 128x1x1 and grid dimensions: 1280x1x1: CUDA_ERROR_ILLEGAL_ADDRESS: an illegal memory access was encountered
	 [[{{node StatefulPartitionedCall}}]] [Op:__inference_train_function_37296]
```
</details> Thank you @AndreasMadsen for the short script to reproduce and for removing the `transformers` dependency to keep the example simpler.

I can reproduce the issue outside docker with CUDA 11.2.0 using the pre-built TF 2.11 pip package. However, the issue disappears when I use CUDA 11.8.0 using the same pip package (the pip package is built with CUDA 11.2 but applications built for one version of CUDA also work for future minor versions of CUDA). So I'm guessing this was a bug with CUDA 11.2 that was fixed somewhere between CUDA 11.2.1 and 11.8.0. Since the TF nightly pip packages are now built with CUDA 11.8, I think this issue is fixed so I'm closing this issue.

It's possible this is a race condition in TF instead of a bug in CUDA, and that updating the CUDA version changes some timings which causes the bug not to occur. But I think this is unlikely, so unless we later see CUDA_ERROR_ILLEGAL_ADDRESS errors in other models, we should consider this issue fixed.

 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/59671">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/59671">No</a>

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

Bug

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

No

### Source

binary

### Tensorflow Version

2.11

### Custom Code

No

### OS Platform and Distribution

Official docker image via apptainer

### Mobile device

_No response_

### Python version

3.8

### Bazel version

_No response_

### GCC/Compiler version

_No response_

### CUDA/cuDNN version

11.2/8.1

### GPU model and memory

Tesla V100-SXM2-32GB

### Current Behaviour?

Using the Official docker image, the provided code results in a CUDA_ERROR_ILLEGAL_ADDRESS. Note, that while the error is fully consistent (happens every time at the first epoch) it is highly sensitive. If a different seed or model-size is used, the error will not reproduce.

I have tried enabling memory growth as suggested elsewhere. However, this only delays the issue to epoch 18.

The error does not reproduce on Google Colab or when using the nightly docker image. However, due to the sensitivity of the issue, it is unclear what exactly this means. The Google Colab issue could be due to different GPU. On my own environment it's a Tesla V100-SXM2-32GB, not a T4.

### Standalone code to reproduce the issue

This uses `transformers == 4.26.0`, and a synthetic dataset file (based on QQP) that is automatically downloaded. Code is available here: https://gist.github.com/AndreasMadsen/2bf669a3cd4c4a8ba964561b9e72279e

Note, I did make an attempt at reproducing it using Colab: https://colab.research.google.com/drive/1zVbNNfz1lZ6xgyoZ7dKWnKzj6poM5XWC?usp=sharing

However, the error does not reproduce on Colab.

### Relevant log output

```shell
2023-02-06 14:07:37.519351: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 AVX512F AVX512_VNNI FMA
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
2023-02-06 14:07:37.759155: I tensorflow/core/util/port.cc:104] 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-02-06 14:07:38.624178: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer.so.7'; dlerror: libnvinfer.so.7: cannot open shared object file: No such file or directory
2023-02-06 14:07:38.624677: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer_plugin.so.7'; dlerror: libnvinfer_plugin.so.7: cannot open shared object file: No such file or directory
2023-02-06 14:07:38.624725: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.
Configuration:
  persistent_dir = /scratch/anmadc/tf_memory_issue
  seed = 0
  batch_size = 16
  model = roberta-base
  jit_compile = True
  precision = mixed_float16

2023-02-06 14:07:40.255950: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 AVX512F AVX512_VNNI FMA
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
2023-02-06 14:07:41.029530: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1613] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 30971 MB memory:  -> device: 0, name: Tesla V100-SXM2-32GB, pci bus id: 0000:af:00.0, compute capability: 7.0
All model checkpoint layers were used when initializing TFRobertaForSequenceClassification.

Some layers of TFRobertaForSequenceClassification were not initialized from the model checkpoint at roberta-base and are newly initialized: ['classifier']
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
WARNING:tensorflow:From /localscratch/anmadc.58777029.0/env/lib/python3.10/site-packages/tensorflow/python/autograph/pyct/static_analysis/liveness.py:83: Analyzer.lamba_check (from tensorflow.python.autograph.pyct.static_analysis.liveness) is deprecated and will be removed after 2023-09-23.
Instructions for updating:
Lambda fuctions will be no more assumed to be used in the statement where they are used, or at least in the same block. https://github.com/tensorflow/tensorflow/issues/56089
2023-02-06 14:08:24.250220: I tensorflow/compiler/xla/service/service.cc:173] XLA service 0x2db1cb30 initialized for platform CUDA (this does not guarantee that XLA will be used). Devices:
2023-02-06 14:08:24.264540: I tensorflow/compiler/xla/service/service.cc:181]   StreamExecutor device (0): Tesla V100-SXM2-32GB, Compute Capability 7.0
2023-02-06 14:08:25.275879: W tensorflow/compiler/tf2xla/kernels/assert_op.cc:38] Ignoring Assert operator tf_roberta_for_sequence_classification/roberta/embeddings/assert_less/Assert/Assert
2023-02-06 14:08:25.309985: I tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.cc:268] disabling MLIR crash reproducer, set env var `MLIR_CRASH_REPRODUCER_DIRECTORY` to enable.
2023-02-06 14:08:25.321306: W tensorflow/compiler/tf2xla/kernels/random_ops.cc:57] Warning: Using tf.random.uniform with XLA compilation will ignore seeds; consider using tf.random.stateless_uniform instead if reproducible behavior is desired. tf_roberta_for_sequence_classification/roberta/embeddings/dropout_36/dropout/random_uniform/RandomUniform
2023-02-06 14:09:41.534987: I tensorflow/compiler/xla/stream_executor/gpu/asm_compiler.cc:325] ptxas warning : Registers are spilled to local memory in function 'sort_19'
ptxas warning : Registers are spilled to local memory in function 'sort_17'

2023-02-06 14:09:41.931415: I tensorflow/compiler/jit/xla_compilation_cache.cc:477] Compiled cluster using XLA!  This line is logged at most once for the lifetime of the process.
2023-02-06 14:09:42.331286: E tensorflow/compiler/xla/stream_executor/cuda/cuda_driver.cc:1159] failed to enqueue async memcpy from device to host: CUDA_ERROR_ILLEGAL_ADDRESS: an illegal memory access was encountered; host dst: 0x2b36f334aeeb; GPU src: 0x2b385f628600; size: 1=0x1
2023-02-06 14:09:42.331472: I tensorflow/compiler/xla/stream_executor/stream.cc:2485] INTERNAL: Unknown error
2023-02-06 14:09:42.331557: I tensorflow/compiler/xla/stream_executor/stream.cc:2489] [stream=0x1b700120,impl=0x52a9880] INTERNAL: stream did not block host until done; was already in an error state
Traceback (most recent call last):
  File "/scratch/anmadc/workspace/economical-roar/reproduce.py", line 150, in <module>
    model.fit(dataset_train_batched, verbose=1, epochs=1)
  File "/localscratch/anmadc.58777029.0/env/lib/python3.10/site-packages/keras/utils/traceback_utils.py", line 70, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/localscratch/anmadc.58777029.0/env/lib/python3.10/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 'StatefulPartitionedCall' defined at (most recent call last):
    File "/scratch/anmadc/workspace/economical-roar/reproduce.py", line 150, in <module>
      model.fit(dataset_train_batched, verbose=1, epochs=1)
    File "/localscratch/anmadc.58777029.0/env/lib/python3.10/site-packages/keras/utils/traceback_utils.py", line 65, in error_handler
      return fn(*args, **kwargs)
    File "/localscratch/anmadc.58777029.0/env/lib/python3.10/site-packages/keras/engine/training.py", line 1650, in fit
      tmp_logs = self.train_function(iterator)
    File "/localscratch/anmadc.58777029.0/env/lib/python3.10/site-packages/keras/engine/training.py", line 1249, in train_function
      return step_function(self, iterator)
    File "/localscratch/anmadc.58777029.0/env/lib/python3.10/site-packages/keras/engine/training.py", line 1233, in step_function
      outputs = model.distribute_strategy.run(run_step, args=(data,))
Node: 'StatefulPartitionedCall'
Failed to retrieve branch_index value on stream 0x1b700120: stream did not block host until done; was already in an error state.
	 [[{{node StatefulPartitionedCall}}]] [Op:__inference_train_function_42578]
```
</details>