import tensorflow as tf

with tf.device("/GPU:0"):
    a = tf.random.normal([1, 2])