Hi,
I am trying to figure out how I can set device to GPU with Tensorflow. With other frameworks I have been able to use GPU with the same instance. I have below code and error. It runs OOM with CPU instantly.
I know with Pytorch I can go as torch.cuda.set_device(), can you help me understand what’s the method in Tensorflow?
tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')
train_encodings = tokenizer(train_texts, truncation=True, padding=True)
train_dataset = tf.data.Dataset.from_tensor_slices((
    dict(train_encodings),
    train_labels
))

model = TFDistilBertForSequenceClassification.from_pretrained('distilbert-base-uncased', num_labels=3)

optimizer = tf.keras.optimizers.Adam(learning_rate=5e-5)
model.compile(optimizer=optimizer, loss= model.compute_loss, metrics=['accuracy'])

model.fit(train_dataset.shuffle(1000).batch(128), epochs=10, batch_size=128,
          validation_data=val_dataset.shuffle(1000).batch(128))

I get error as below:
2021-11-08 05:50:31.398734: W tensorflow/core/framework/op_kernel.cc:1745] OP_REQUIRES failed at transpose_op.cc:183 : RESOURCE_EXHAUSTED: OOM when allocating tensor with shape[128,512,12,64] and type float on /job:localhost/replica:0/task:0/device:CPU:0 by allocator cpu