fix: getting device of KerasVariable

ivy-llc · Jul 19, 2024 · 3805db6 · 3805db6
1 parent ce5f797
commit 3805db6
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/ivy/functional/backends/tensorflow/device.py b/ivy/functional/backends/tensorflow/device.py
@@ -28,6 +28,9 @@ def dev(
     *,
     as_native: bool = False,
 ) -> Union[ivy.Device, str]:
+    if "keras.src.backend.tensorflow.core.Variable" in str(x.__class__):
+        # Read the underlying tensor of a KerasVariable to get the device
+        x = x.value
     if isinstance(x, tf.TensorArray):
         # Read the underlying tensor being wrapped to get the device.
         x = x.stack()