`Could not find valid device for node.`

[DarrenZhang01]

As titled

stax.py", line 3111, in _conv_kernel_diagonal_spatial
    mat = np.reshape(mat, (-1,) + mat[batch_ndim:])
  File "/Library/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/trax/tf_numpy/numpy_impl/array_ops.py", line 858, in reshape
    r = tf.reshape(a.data, newshape)
  File "/Library/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/tensorflow/python/ops/array_ops.py", line 193, in reshape
    result = gen_array_ops.reshape(tensor, shape, name)
  File "/Library/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/tensorflow/python/ops/gen_array_ops.py", line 8084, in reshape
    _ops.raise_from_not_ok_status(e, name)
  File "/Library/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/tensorflow/python/framework/ops.py", line 6653, in raise_from_not_ok_status
    six.raise_from(core._status_to_exception(e.code, message), None)
  File "<string>", line 3, in raise_from
tensorflow.python.framework.errors_impl.NotFoundError: Could not find valid device for node.
Node:{{node Reshape}}
All kernels registered for op Reshape :
  device='XLA_CPU'; Tshape in [DT_INT32, DT_INT64]; T in [DT_FLOAT, DT_DOUBLE, DT_INT32, DT_UINT8, DT_INT16, ..., DT_UINT16, DT_COMPLEX128, DT_HALF, DT_UINT32, DT_UINT64]
  device='XLA_CPU_JIT'; Tshape in [DT_INT32, DT_INT64]; T in [DT_FLOAT, DT_DOUBLE, DT_INT32, DT_UINT8, DT_INT16, ..., DT_UINT16, DT_COMPLEX128, DT_HALF, DT_UINT32, DT_UINT64]
  device='GPU'; T in [DT_HALF]; Tshape in [DT_INT32]
  device='GPU'; T in [DT_HALF]; Tshape in [DT_INT64]
  device='GPU'; T in [DT_BFLOAT16]; Tshape in [DT_INT32]
  device='GPU'; T in [DT_BFLOAT16]; Tshape in [DT_INT64]
  device='GPU'; T in [DT_FLOAT]; Tshape in [DT_INT32]
  device='GPU'; T in [DT_FLOAT]; Tshape in [DT_INT64]
  device='GPU'; T in [DT_DOUBLE]; Tshape in [DT_INT32]
  device='GPU'; T in [DT_DOUBLE]; Tshape in [DT_INT64]
  device='GPU'; T in [DT_INT64]; Tshape in [DT_INT32]
  device='GPU'; T in [DT_INT64]; Tshape in [DT_INT64]
  device='GPU'; T in [DT_UINT16]; Tshape in [DT_INT32]
  device='GPU'; T in [DT_UINT16]; Tshape in [DT_INT64]
  device='GPU'; T in [DT_INT16]; Tshape in [DT_INT32]
  device='GPU'; T in [DT_INT16]; Tshape in [DT_INT64]
  device='GPU'; T in [DT_UINT8]; Tshape in [DT_INT32]
  device='GPU'; T in [DT_UINT8]; Tshape in [DT_INT64]
  device='GPU'; T in [DT_INT8]; Tshape in [DT_INT32]
  device='GPU'; T in [DT_INT8]; Tshape in [DT_INT64]
  device='GPU'; T in [DT_COMPLEX64]; Tshape in [DT_INT32]
  device='GPU'; T in [DT_COMPLEX64]; Tshape in [DT_INT64]
  device='GPU'; T in [DT_COMPLEX128]; Tshape in [DT_INT32]
  device='GPU'; T in [DT_COMPLEX128]; Tshape in [DT_INT64]
  device='GPU'; T in [DT_BOOL]; Tshape in [DT_INT32]
  device='GPU'; T in [DT_BOOL]; Tshape in [DT_INT64]
  device='CPU'; Tshape in [DT_INT32]
  device='CPU'; Tshape in [DT_INT64]
 [Op:Reshape]

[DarrenZhang01]

Solved in 6af6710