How to enable set device as GPU with Tensorflow?


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 =

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']), epochs=10, batch_size=128,

I get error as below:

2021-11-08 05:50:31.398734: W tensorflow/core/framework/] OP_REQUIRES failed at : 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

Normally, Keras automatically uses the GPU if it’s available.

You can check the available GPUs as follows:

import tensorflow as tf

print("Num GPUs Available: ", len(tf.config.list_physical_devices('GPU')))

Hey thanks for that. Seems that tensorflow actually doesn’t register my GPU when I run


I get error as below:

2021-11-08 22:08:10.636081: I tensorflow/stream_executor/cuda/] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2021-11-08 22:08:10.636549: W tensorflow/core/common_runtime/gpu/] Cannot dlopen some GPU libraries. Please make sure the missing libraries mentioned above are installed properly if you would like to use GPU. Follow the guide at for how to download and setup the required libraries for your platform.
Skipping registering GPU devices...

I can use the GPU with Pytorch definitely. When I change the code to Pytorch and place the model on GPU, and start training, I can see GPU usage increasing from nvidia-smi.

I’m using Tensorflow 2.7.0 as well, so it’s not any of the older versions where you need tensorflow-gpu as well.

I’m using cuda V11.0.221 and tried reinstalling but it seems that my account doesn’t have the permission to do that unfortunately.

Any ideas to tackle this would be appreciated highly :slight_smile: