ut train fails on GPU with JIT compilation failed

[Shubhamcl]

Using ubuntu 22, training works fine on CPU but when --num_gpus=1 I get this error stack.

This stack appears on following the instructions for the demo.

I first thought it is a tensorflow issue so I ran training on GPU using example from tensorflow tutorials, but that worked fine.

Detected at node 'SelectV2' defined at (most recent call last):
File "/home/shubham/anaconda3/envs/u-sleep/lib/python3.9/threading.py", line 908, in _bootstrap
self._bootstrap_inner()
File "/home/shubham/anaconda3/envs/u-sleep/lib/python3.9/threading.py", line 950, in _bootstrap_inner
self.run()
File "/home/shubham/anaconda3/envs/u-sleep/lib/python3.9/site-packages/keras/engine/training.py", line 1000, in run_step
outputs = model.train_step(data)
File "/home/shubham/anaconda3/envs/u-sleep/lib/python3.9/site-packages/keras/engine/training.py", line 864, in train_step
return self.compute_metrics(x, y, y_pred, sample_weight)
File "/home/shubham/anaconda3/envs/u-sleep/lib/python3.9/site-packages/keras/engine/training.py", line 957, in compute_metrics
self.compiled_metrics.update_state(y, y_pred, sample_weight)
File "/home/shubham/anaconda3/envs/u-sleep/lib/python3.9/site-packages/keras/engine/compile_utils.py", line 459, in update_state
metric_obj.update_state(y_t, y_p, sample_weight=mask)
File "/home/shubham/anaconda3/envs/u-sleep/lib/python3.9/site-packages/utime/evaluation/utils.py", line 22, in wrapper
mask = tf.where(tf.logical_and(
Node: 'SelectV2'
Detected at node 'SelectV2' defined at (most recent call last):
File "/home/shubham/anaconda3/envs/u-sleep/lib/python3.9/threading.py", line 908, in _bootstrap
self._bootstrap_inner()
File "/home/shubham/anaconda3/envs/u-sleep/lib/python3.9/threading.py", line 950, in _bootstrap_inner
self.run()
File "/home/shubham/anaconda3/envs/u-sleep/lib/python3.9/site-packages/keras/engine/training.py", line 1000, in run_step
outputs = model.train_step(data)
File "/home/shubham/anaconda3/envs/u-sleep/lib/python3.9/site-packages/keras/engine/training.py", line 864, in train_step
return self.compute_metrics(x, y, y_pred, sample_weight)
File "/home/shubham/anaconda3/envs/u-sleep/lib/python3.9/site-packages/keras/engine/training.py", line 957, in compute_metrics
self.compiled_metrics.update_state(y, y_pred, sample_weight)
File "/home/shubham/anaconda3/envs/u-sleep/lib/python3.9/site-packages/keras/engine/compile_utils.py", line 459, in update_state
metric_obj.update_state(y_t, y_p, sample_weight=mask)
File "/home/shubham/anaconda3/envs/u-sleep/lib/python3.9/site-packages/utime/evaluation/utils.py", line 22, in wrapper
mask = tf.where(tf.logical_and(
Node: 'SelectV2'
2 root error(s) found.
(0) UNKNOWN: JIT compilation failed.
[[{{node SelectV2}}]]
[[div_no_nan_1/ReadVariableOp/_12]]
(1) UNKNOWN: JIT compilation failed.
[[{{node SelectV2}}]]
0 successful operations.
0 derived errors ignored. [Op:__inference_train_function_13068]