Skip to content

Commit fc323e1

Browse files
authored
util: Explicitly call __dlpack__ built-in method in xp2tensorflow (#757)
`tf.experimental.dlpack.from_dlpack` expects a `PyCapsule` object.
1 parent fba3bf0 commit fc323e1

File tree

1 file changed

+2
-1
lines changed

1 file changed

+2
-1
lines changed

thinc/util.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -404,7 +404,8 @@ def xp2tensorflow(
404404
dlpack_tensor = xp_tensor.toDlpack() # type: ignore
405405
tf_tensor = tf.experimental.dlpack.from_dlpack(dlpack_tensor)
406406
elif hasattr(xp_tensor, "__dlpack__"):
407-
tf_tensor = tf.experimental.dlpack.from_dlpack(xp_tensor)
407+
dlpack_tensor = xp_tensor.__dlpack__() # type: ignore
408+
tf_tensor = tf.experimental.dlpack.from_dlpack(dlpack_tensor)
408409
else:
409410
tf_tensor = tf.convert_to_tensor(xp_tensor)
410411
if as_variable:

0 commit comments

Comments
 (0)