Keras 3.3.3: Tensor on device cuda:0 is not on the expected device meta

[jeffDebug]

I encountered this error when I was running the training, with the backend set to torch
"RuntimeError: Tensor on device cuda:0 is not on the expected device meta"

The more stack trace messages are as followings.

  File "C:\jeff\conda\envs\keras3\lib\site-packages\torch\utils\_device.py", line 78, in __torch_function__
    return func(*args, **kwargs)
  File "C:\jeff\conda\envs\keras3\lib\site-packages\torch\_prims_common\wrappers.py", line 252, in _fn
    result = fn(*args, **kwargs)
  File "C:\jeff\conda\envs\keras3\lib\site-packages\torch\_prims_common\wrappers.py", line 137, in _fn
    result = fn(**bound.arguments)
  File "C:\jeff\conda\envs\keras3\lib\site-packages\torch\_refs\__init__.py", line 1091, in add
    output = prims.add(a, b)
  File "C:\jeff\conda\envs\keras3\lib\site-packages\torch\_ops.py", line 594, in __call__
    return self_._op(*args, **kwargs)
  File "C:\jeff\conda\envs\keras3\lib\site-packages\torch\_prims\__init__.py", line 359, in _prim_elementwise_meta
    utils.check_same_device(*args_, allow_cpu_scalar_tensors=True)
  File "C:\jeff\conda\envs\keras3\lib\site-packages\torch\_prims_common\__init__.py", line 740, in check_same_device
    raise RuntimeError(msg)
RuntimeError: Tensor on device cuda:0 is not on the expected device meta!

I have installed:
keras==3.3.3
torch==2.3.0+cu118

Please help or any hint will be appreciated.

[ghsanti]

Could you try and share some minimal example in Collab, so that we can reproduce it?

[jeffDebug]

I have found out why. It's due to some of our legacy code, where some tensorflow (2.16) API are still called.

By the way, thanks for the response. Also my apologies for that there is no easy way for me to share some minimum example to reproduce this.