Repository: tensorflow
Issue: could not run GPU on jupyter 
Labels: type:bug,type:feature,comp:gpu,TF 2.12

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

Bug

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

Yes

### Source

source

### Tensorflow Version

2.12.0

### Custom Code

Yes

### OS Platform and Distribution

_No response_

### Mobile device

_No response_

### Python version

3.10.11

### Bazel version

_No response_

### GCC/Compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current Behaviour?

I could not get Tensorflow to run on GPUs.

TF sees the GPUs on terminal, but not on jupyter lab.

**Edited**
Found a solution to see it on jupyterlab, but must manually repeat . Still, erratic misconfiguration seems to happen.


### Standalone code to reproduce the issue

```shell
I could not get Tensorflow to run on GPUs.

TF sees the GPUs on terminal, but not on jupyter lab.

** edited **


Eventually, after hours, I found a temporary solution in setting the paths each time, before I launch `jupyter lab` :



CUDNN_PATH=$(dirname $(python -c "import nvidia.cudnn;print(nvidia.cudnn.__file__)"))
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$CONDA_PREFIX/lib/:$CUDNN_PATH/lib

```

NB : If I was to set 
```
mkdir -p $CONDA_PREFIX/etc/conda/activate.d
echo 'CUDNN_PATH=$(dirname $(python -c "import nvidia.cudnn;print(nvidia.cudnn.__file__)"))' >> $CONDA_PREFIX/etc/conda/activate.d/env_vars.sh
echo 'export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$CONDA_PREFIX/lib/:$CUDNN_PATH/lib' >> $CONDA_PREFIX/etc/conda/activate.d/env_vars.sh
```

as in `https://www.tensorflow.org/install/pip`

Somehow it messy with jupyter, bcause the selected kernel from jupyter would not correspond to the kernel set from the script.

Now I can see the GPUs on jupyter, but still - it crashes! And before, on CPU, it was not.

When I run this script, using this library :

https://pypi.org/project/umap-learn/

installed as in the description:

```
embedder = ParametricUMAP(
    ## all the params ...
)

# now launch on GPU
with tf.device('/GPU:0'):
    embedding =  embedder.fit_transform(np.array([t.ravel() for t in train_data]))
```

the code fails with the log output below.

If I close the jupyter lab connection, go back on the conda environment, 
set again:

```
CUDNN_PATH=$(dirname $(python -c "import nvidia.cudnn;print(nvidia.cudnn.__file__)"))
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$CONDA_PREFIX/lib/:$CUDNN_PATH/lib
```
I
and 

```
# Install NVCC
conda install -c nvidia cuda-nvcc=11.3.58
# Configure the XLA cuda directory
mkdir -p $CONDA_PREFIX/etc/conda/activate.d
printf 'export XLA_FLAGS=--xla_gpu_cuda_data_dir=$CONDA_PREFIX/lib/\n' >> $CONDA_PREFIX/etc/conda/activate.d/env_vars.sh
source $CONDA_PREFIX/etc/conda/activate.d/env_vars.sh
# Copy libdevice file to the required path
mkdir -p $CONDA_PREFIX/lib/nvvm/libdevice
cp $CONDA_PREFIX/lib/libdevice.10.bc $CONDA_PREFIX/lib/nvvm/libdevice/

```

And relaunch jupterlab, this time I get the GPU seen also in jupyter lab.

Running 

```
with tf.device('/GPU:0'):
    spec_embedding = spec_embedder.fit_transform(np.array([t.ravel() for t in train_data]))

```

It yields

```
2023-06-06 21:55:04.944257: 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 'gradients/split_grad/concat/split/split_dim' with dtype int32
	 [[{{node gradients/split_grad/concat/split/split_dim}}]]
```

And cannot understand why it raise warning or error with `CPU` device, if I run on `GPU`. 

Not even sure it falled back to CPU, actually.


---

Please advice how to sync jupter lab and conda.
I did follow up with `ikernel` but it seems, after a lot of checking, that environment variables are not correctly read. Not sure if `kernel.json` fails to be updated properly, with thepath of the cuda libraries.

Please consider add a guide on:
https://www.tensorflow.org/install/pip

for running TF on jupyter.

My situation is that I need to run from a remote cluster, and I think it is a frequent situation.
hope this feedback is useful.
```


### Relevant log output

```shell
2023-06-06 21:40:08.385763: 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 'gradients/split_grad/concat/split/split_dim' with dtype int32
	 [[{{node gradients/split_grad/concat/split/split_dim}}]]
2023-06-06 21:40:10.938429: I tensorflow/compiler/xla/stream_executor/cuda/cuda_dnn.cc:424] Loaded cuDNN version 8600
2023-06-06 21:40:11.641291: I tensorflow/compiler/xla/service/service.cc:169] XLA service 0x7f89f489b8d0 initialized for platform CUDA (this does not guarantee that XLA will be used). Devices:
2023-06-06 21:40:11.641337: I tensorflow/compiler/xla/service/service.cc:177]   StreamExecutor device (0): NVIDIA GeForce RTX 2080 Ti, Compute Capability 7.5
2023-06-06 21:40:11.641346: I tensorflow/compiler/xla/service/service.cc:177]   StreamExecutor device (1): NVIDIA GeForce RTX 2080 Ti, Compute Capability 7.5
2023-06-06 21:40:11.641353: I tensorflow/compiler/xla/service/service.cc:177]   StreamExecutor device (2): NVIDIA GeForce RTX 2080 Ti, Compute Capability 7.5
2023-06-06 21:40:11.641359: I tensorflow/compiler/xla/service/service.cc:177]   StreamExecutor device (3): NVIDIA GeForce RTX 2080 Ti, Compute Capability 7.5
2023-06-06 21:40:11.646324: I tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.cc:269] disabling MLIR crash reproducer, set env var `MLIR_CRASH_REPRODUCER_DIRECTORY` to enable.
2023-06-06 21:40:11.673008: W tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc:530] Can't find libdevice directory ${CUDA_DIR}/nvvm/libdevice. This may result in compilation or runtime failures, if the program we try to run uses routines from libdevice.
Searched for CUDA in the following directories:
  ./cuda_sdk_lib
  /usr/local/cuda-11.8
  /usr/local/cuda
  .
You can choose the search directory by setting xla_gpu_cuda_data_dir in HloModule's DebugOptions.  For most apps, setting the environment variable XLA_FLAGS=--xla_gpu_cuda_data_dir=/path/to/cuda will work.
2023-06-06 21:40:11.673259: W tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc:274] libdevice is required by this HLO module but was not found at ./libdevice.10.bc
2023-06-06 21:40:11.673658: W tensorflow/core/framework/op_kernel.cc:1830] OP_REQUIRES failed at xla_ops.cc:362 : INTERNAL: libdevice not found at ./libdevice.10.bc
2023-06-06 21:40:11.673685: 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: libdevice not found at ./libdevice.10.bc
	 [[{{node StatefulPartitionedCall_16}}]]
2023-06-06 21:40:11.700045: W tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc:274] libdevice is required by this HLO module but was not found at ./libdevice.10.bc
2023-06-06 21:40:11.700414: W tensorflow/core/framework/op_kernel.cc:1830] OP_REQUIRES failed at xla_ops.cc:362 : INTERNAL: libdevice not found at ./libdevice.10.bc
2023-06-06 21:40:11.729242: W tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc:274] libdevice is required by this HLO module but was not found at ./libdevice.10.bc
2023-06-06 21:40:11.729590: W tensorflow/core/framework/op_kernel.cc:1830] OP_REQUIRES failed at xla_ops.cc:362 : INTERNAL: libdevice not found at ./libdevice.10.bc
2023-06-06 21:40:11.755959: W tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc:274] libdevice is required by this HLO module but was not found at ./libdevice.10.bc
2023-06-06 21:40:11.756308: W tensorflow/core/framework/op_kernel.cc:1830] OP_REQUIRES failed at xla_ops.cc:362 : INTERNAL: libdevice not found at ./libdevice.10.bc
2023-06-06 21:40:11.783041: W tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc:274] libdevice is required by this HLO module but was not found at ./libdevice.10.bc
2023-06-06 21:40:11.783397: W tensorflow/core/framework/op_kernel.cc:1830] OP_REQUIRES failed at xla_ops.cc:362 : INTERNAL: libdevice not found at ./libdevice.10.bc
2023-06-06 21:40:11.809786: W tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc:274] libdevice is required by this HLO module but was not found at ./libdevice.10.bc
2023-06-06 21:40:11.810134: W tensorflow/core/framework/op_kernel.cc:1830] OP_REQUIRES failed at xla_ops.cc:362 : INTERNAL: libdevice not found at ./libdevice.10.bc
2023-06-06 21:40:11.836777: W tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc:274] libdevice is required by this HLO module but was not found at ./libdevice.10.bc
2023-06-06 21:40:11.837129: W tensorflow/core/framework/op_kernel.cc:1830] OP_REQUIRES failed at xla_ops.cc:362 : INTERNAL: libdevice not found at ./libdevice.10.bc
2023-06-06 21:40:11.864411: W tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc:274] libdevice is required by this HLO module but was not found at ./libdevice.10.bc
2023-06-06 21:40:11.864767: W tensorflow/core/framework/op_kernel.cc:1830] OP_REQUIRES failed at xla_ops.cc:362 : INTERNAL: libdevice not found at ./libdevice.10.bc
2023-06-06 21:40:11.892388: W tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc:274] libdevice is required by this HLO module but was not found at ./libdevice.10.bc
2023-06-06 21:40:11.892738: W tensorflow/core/framework/op_kernel.cc:1830] OP_REQUIRES failed at xla_ops.cc:362 : INTERNAL: libdevice not found at ./libdevice.10.bc
2023-06-06 21:40:11.919296: W tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc:274] libdevice is required by this HLO module but was not found at ./libdevice.10.bc
2023-06-06 21:40:11.919647: W tensorflow/core/framework/op_kernel.cc:1830] OP_REQUIRES failed at xla_ops.cc:362 : INTERNAL: libdevice not found at ./libdevice.10.bc
2023-06-06 21:40:11.946506: W tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc:274] libdevice is required by this HLO module but was not found at ./libdevice.10.bc
2023-06-06 21:40:11.946866: W tensorflow/core/framework/op_kernel.cc:1830] OP_REQUIRES failed at xla_ops.cc:362 : INTERNAL: libdevice not found at ./libdevice.10.bc
2023-06-06 21:40:11.974287: W tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc:274] libdevice is required by this HLO module but was not found at ./libdevice.10.bc
2023-06-06 21:40:11.974699: W tensorflow/core/framework/op_kernel.cc:1830] OP_REQUIRES failed at xla_ops.cc:362 : INTERNAL: libdevice not found at ./libdevice.10.bc
2023-06-06 21:40:12.245782: W tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc:274] libdevice is required by this HLO module but was not found at ./libdevice.10.bc
2023-06-06 21:40:12.246188: W tensorflow/core/framework/op_kernel.cc:1830] OP_REQUIRES failed at xla_ops.cc:362 : INTERNAL: libdevice not found at ./libdevice.10.bc
2023-06-06 21:40:12.272303: W tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc:274] libdevice is required by this HLO module but was not found at ./libdevice.10.bc
2023-06-06 21:40:12.272657: W tensorflow/core/framework/op_kernel.cc:1830] OP_REQUIRES failed at xla_ops.cc:362 : INTERNAL: libdevice not found at ./libdevice.10.bc
2023-06-06 21:40:12.322559: W tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc:274] libdevice is required by this HLO module but was not found at ./libdevice.10.bc
2023-06-06 21:40:12.322899: W tensorflow/core/framework/op_kernel.cc:1830] OP_REQUIRES failed at xla_ops.cc:362 : INTERNAL: libdevice not found at ./libdevice.10.bc
2023-06-06 21:40:12.350255: W tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc:274] libdevice is required by this HLO module but was not found at ./libdevice.10.bc
2023-06-06 21:40:12.350736: W tensorflow/core/framework/op_kernel.cc:1830] OP_REQUIRES failed at xla_ops.cc:362 : INTERNAL: libdevice not found at ./libdevice.10.bc
2023-06-06 21:40:12.379576: W tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc:274] libdevice is required by this HLO module but was not found at ./libdevice.10.bc
2023-06-06 21:40:12.379966: W tensorflow/core/framework/op_kernel.cc:1830] OP_REQUIRES failed at xla_ops.cc:362 : INTERNAL: libdevice not found at ./libdevice.10.bc
2023-06-06 21:40:12.406780: W tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc:274] libdevice is required by this HLO module but was not found at ./libdevice.10.bc
2023-06-06 21:40:12.407233: W tensorflow/core/framework/op_kernel.cc:1830] OP_REQUIRES failed at xla_ops.cc:362 : INTERNAL: libdevice not found at ./libdevice.10.bc

---------------------------------------------------------------------------
InternalError                             Traceback (most recent call last)
Cell In[33], line 31
      7 spec_embedder = ParametricUMAP(
      8     metric = 'euclidean',
      9     min_dist = 0.1, 
   (...)
     27     n_training_epochs=1
     28 )
     30 with tf.device('/GPU:0'):
---> 31     spec_embedding = spec_embedder.fit_transform(np.array([t.ravel() for t in train_data]))

File ~/miniconda3/envs/conda01/lib/python3.10/site-packages/umap/parametric_umap.py:217, in ParametricUMAP.fit_transform(self, X, y, precomputed_distances)
    215     return super().fit_transform(precomputed_distances, y)
    216 else:
--> 217     return super().fit_transform(X, y)

File ~/miniconda3/envs/conda01/lib/python3.10/site-packages/umap/umap_.py:2772, in UMAP.fit_transform(self, X, y)
   2742 def fit_transform(self, X, y=None):
   2743     """Fit X into an embedded space and return that transformed
   2744     output.
   2745 
   (...)
   2770         Local radii of data points in the embedding (log-transformed).
   2771     """
-> 2772     self.fit(X, y)
   2773     if self.transform_mode == "embedding":
   2774         if self.output_dens:

File ~/miniconda3/envs/conda01/lib/python3.10/site-packages/umap/parametric_umap.py:202, in ParametricUMAP.fit(self, X, y, precomputed_distances)
    200     return super().fit(precomputed_distances, y)
    201 else:
--> 202     return super().fit(X, y)

File ~/miniconda3/envs/conda01/lib/python3.10/site-packages/umap/umap_.py:2684, in UMAP.fit(self, X, y)
   2681     print(ts(), "Construct embedding")
   2683 if self.transform_mode == "embedding":
-> 2684     self.embedding_, aux_data = self._fit_embed_data(
   2685         self._raw_data[index],
   2686         self.n_epochs,
   2687         init,
   2688         random_state,  # JH why raw data?
   2689     )
   2690     # Assign any points that are fully disconnected from our manifold(s) to have embedding
   2691     # coordinates of np.nan.  These will be filtered by our plotting functions automatically.
   2692     # They also prevent users from being deceived a distance query to one of these points.
   2693     # Might be worth moving this into simplicial_set_embedding or _fit_embed_data
   2694     disconnected_vertices = np.array(self.graph_.sum(axis=1)).flatten() == 0

File ~/miniconda3/envs/conda01/lib/python3.10/site-packages/umap/parametric_umap.py:462, in ParametricUMAP._fit_embed_data(self, X, n_epochs, init, random_state)
    459     validation_data = None
    461 # create embedding
--> 462 history = self.parametric_model.fit(
    463     edge_dataset,
    464     epochs=self.loss_report_frequency * self.n_training_epochs,
    465     steps_per_epoch=steps_per_epoch,
    466     max_queue_size=100,
    467     validation_data=validation_data,
    468     **self.keras_fit_kwargs
    469 )
    470 # save loss history dictionary
    471 self._history = history.history

File ~/.local/lib/python3.10/site-packages/keras/utils/traceback_utils.py:70, in filter_traceback.<locals>.error_handler(*args, **kwargs)
     67     filtered_tb = _process_traceback_frames(e.__traceback__)
     68     # To get the full stack trace, call:
     69     # `tf.debugging.disable_traceback_filtering()`
---> 70     raise e.with_traceback(filtered_tb) from None
     71 finally:
     72     del filtered_tb

File ~/miniconda3/envs/conda01/lib/python3.10/site-packages/tensorflow/python/eager/execute.py:52, in quick_execute(op_name, num_outputs, inputs, attrs, ctx, name)
     50 try:
     51   ctx.ensure_initialized()
---> 52   tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
     53                                       inputs, attrs, num_outputs)
     54 except core._NotOkStatusException as e:
     55   if name is not None:

InternalError: Graph execution error:

Detected at node 'StatefulPartitionedCall_16' defined at (most recent call last):
    File "/home/h21/luas6629/miniconda3/envs/conda01/lib/python3.10/runpy.py", line 196, in _run_module_as_main
      return _run_code(code, main_globals, None,
    File "/home/h21/luas6629/miniconda3/envs/conda01/lib/python3.10/runpy.py", line 86, in _run_code
      exec(code, run_globals)
    File "/home/h21/luas6629/miniconda3/envs/conda01/lib/python3.10/site-packages/ipykernel_launcher.py", line 17, in <module>
      app.launch_new_instance()
    File "/home/h21/luas6629/miniconda3/envs/conda01/lib/python3.10/site-packages/traitlets/config/application.py", line 992, in launch_instance
      app.start()
    File "/home/h21/luas6629/miniconda3/envs/conda01/lib/python3.10/site-packages/ipykernel/kernelapp.py", line 711, in start
      self.io_loop.start()
    File "/home/h21/luas6629/miniconda3/envs/conda01/lib/python3.10/site-packages/tornado/platform/asyncio.py", line 215, in start
      self.asyncio_loop.run_forever()
    File "/home/h21/luas6629/miniconda3/envs/conda01/lib/python3.10/asyncio/base_events.py", line 603, in run_forever
      self._run_once()
    File "/home/h21/luas6629/miniconda3/envs/conda01/lib/python3.10/asyncio/base_events.py", line 1909, in _run_once
      handle._run()
    File "/home/h21/luas6629/miniconda3/envs/conda01/lib/python3.10/asyncio/events.py", line 80, in _run
      self._context.run(self._callback, *self._args)
    File "/home/h21/luas6629/miniconda3/envs/conda01/lib/python3.10/site-packages/ipykernel/kernelbase.py", line 510, in dispatch_queue
      await self.process_one()
    File "/home/h21/luas6629/miniconda3/envs/conda01/lib/python3.10/site-packages/ipykernel/kernelbase.py", line 499, in process_one
      await dispatch(*args)
    File "/home/h21/luas6629/miniconda3/envs/conda01/lib/python3.10/site-packages/ipykernel/kernelbase.py", line 406, in dispatch_shell
      await result
    File "/home/h21/luas6629/miniconda3/envs/conda01/lib/python3.10/site-packages/ipykernel/kernelbase.py", line 729, in execute_request
      reply_content = await reply_content
    File "/home/h21/luas6629/miniconda3/envs/conda01/lib/python3.10/site-packages/ipykernel/ipkernel.py", line 411, in do_execute
      res = shell.run_cell(
    File "/home/h21/luas6629/miniconda3/envs/conda01/lib/python3.10/site-packages/ipykernel/zmqshell.py", line 531, in run_cell
      return super().run_cell(*args, **kwargs)
    File "/home/h21/luas6629/miniconda3/envs/conda01/lib/python3.10/site-packages/IPython/core/interactiveshell.py", line 3006, in run_cell
      result = self._run_cell(
    File "/home/h21/luas6629/miniconda3/envs/conda01/lib/python3.10/site-packages/IPython/core/interactiveshell.py", line 3061, in _run_cell
      result = runner(coro)
    File "/home/h21/luas6629/miniconda3/envs/conda01/lib/python3.10/site-packages/IPython/core/async_helpers.py", line 129, in _pseudo_sync_runner
      coro.send(None)
    File "/home/h21/luas6629/miniconda3/envs/conda01/lib/python3.10/site-packages/IPython/core/interactiveshell.py", line 3266, in run_cell_async
      has_raised = await self.run_ast_nodes(code_ast.body, cell_name,
    File "/home/h21/luas6629/miniconda3/envs/conda01/lib/python3.10/site-packages/IPython/core/interactiveshell.py", line 3445, in run_ast_nodes
      if await self.run_code(code, result, async_=asy):
    File "/home/h21/luas6629/miniconda3/envs/conda01/lib/python3.10/site-packages/IPython/core/interactiveshell.py", line 3505, in run_code
      exec(code_obj, self.user_global_ns, self.user_ns)
    File "/tmp/ipykernel_3110734/3903349550.py", line 31, in <module>
      spec_embedding = spec_embedder.fit_transform(np.array([t.ravel() for t in train_data]))
    File "/home/h21/luas6629/miniconda3/envs/conda01/lib/python3.10/site-packages/umap/parametric_umap.py", line 217, in fit_transform
      return super().fit_transform(X, y)
    File "/home/h21/luas6629/miniconda3/envs/conda01/lib/python3.10/site-packages/umap/umap_.py", line 2772, in fit_transform
      self.fit(X, y)
    File "/home/h21/luas6629/miniconda3/envs/conda01/lib/python3.10/site-packages/umap/parametric_umap.py", line 202, in fit
      return super().fit(X, y)
    File "/home/h21/luas6629/miniconda3/envs/conda01/lib/python3.10/site-packages/umap/umap_.py", line 2684, in fit
      self.embedding_, aux_data = self._fit_embed_data(
    File "/home/h21/luas6629/miniconda3/envs/conda01/lib/python3.10/site-packages/umap/parametric_umap.py", line 462, in _fit_embed_data
      history = self.parametric_model.fit(
    File "/home/h21/luas6629/.local/lib/python3.10/site-packages/keras/utils/traceback_utils.py", line 65, in error_handler
      return fn(*args, **kwargs)
    File "/home/h21/luas6629/.local/lib/python3.10/site-packages/keras/engine/training.py", line 1685, in fit
      tmp_logs = self.train_function(iterator)
    File "/home/h21/luas6629/.local/lib/python3.10/site-packages/keras/engine/training.py", line 1284, in train_function
      return step_function(self, iterator)
    File "/home/h21/luas6629/.local/lib/python3.10/site-packages/keras/engine/training.py", line 1268, in step_function
      outputs = model.distribute_strategy.run(run_step, args=(data,))
    File "/home/h21/luas6629/.local/lib/python3.10/site-packages/keras/engine/training.py", line 1249, in run_step
      outputs = model.train_step(data)
    File "/home/h21/luas6629/miniconda3/envs/conda01/lib/python3.10/site-packages/umap/parametric_umap.py", line 1150, in train_step
      self.optimizer.apply_gradients(zip(gradients, trainable_vars))
    File "/home/h21/luas6629/.local/lib/python3.10/site-packages/keras/optimizers/optimizer.py", line 1174, in apply_gradients
      return super().apply_gradients(grads_and_vars, name=name)
    File "/home/h21/luas6629/.local/lib/python3.10/site-packages/keras/optimizers/optimizer.py", line 650, in apply_gradients
      iteration = self._internal_apply_gradients(grads_and_vars)
    File "/home/h21/luas6629/.local/lib/python3.10/site-packages/keras/optimizers/optimizer.py", line 1200, in _internal_apply_gradients
      return tf.__internal__.distribute.interim.maybe_merge_call(
    File "/home/h21/luas6629/.local/lib/python3.10/site-packages/keras/optimizers/optimizer.py", line 1250, in _distributed_apply_gradients_fn
      distribution.extended.update(
    File "/home/h21/luas6629/.local/lib/python3.10/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: 'StatefulPartitionedCall_16'
libdevice not found at ./libdevice.10.bc
	 [[{{node StatefulPartitionedCall_16}}]] [Op:__inference_train_function_4496]
```
</details>