message: add test for gpu
diff: @@ -25,6 +25,7 @@ from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import ops
  from tensorflow.python.platform import test
  +from tensorflow.python.ops import array_ops
  int_dtypes = [
  np.int8, np.int16, np.int32, np.int64, np.uint8, np.uint16, np.uint32,
  @@ -59,11 +60,19 @@ class DLPackTest(parameterized.TestCase, test.TestCase):
  def testRoundTrip(self, dtype, shape):
  np.random.seed(42)
  np_array = np.random.randint(0, 10, shape)
  - tf_tensor = constant_op.constant(np_array, dtype=dtype)
  + # copy to gpu if available
  + tf_tensor = array_ops.identity(constant_op.constant(np_array, dtype=dtype))
  + tf_tensor_device = tf_tensor.device
  + tf_tensor_dtype = tf_tensor.dtype
  dlcapsule = dlpack.to_dlpack(tf_tensor)
  del tf_tensor # should still work
  tf_tensor2 = dlpack.from_dlpack(dlcapsule)
  self.assertAllClose(np_array, tf_tensor2)
  + if tf_tensor_dtype == dtypes.int32:
  + # int32 tensor always on cpu for now
  + self.assertEqual(tf_tensor2.device, "/job:localhost/replica:0/task:0/device:CPU:0")
  + else:
  + self.assertEqual(tf_tensor_device, tf_tensor2.device)
  def testTensorsCanBeConsumedOnceOnly(self):
  np.random.seed(42)
  