Does Pytorch automaticly use the GPU

Question

I'm new to using torch and calculate on huge amount of data. I want to create embeddings for an large text corpus and write my embedding function, it works well but it seems pretty slow, so I'm not sure if it really uses the GPU or does still uses the CPU.

import torch
from transformers import BartTokenizer, BartForConditionalGeneration
model_name = 'facebook/bart-base'
def load_model():
 # load the pretrained bart model
 model = BartForConditionalGeneration.from_pretrained(model_name)
 # load the tokenizer
 tokenizer = BartTokenizer.from_pretrained(model_name)
 return model, tokenizer
def calculate_text_embeddings(text, model, tokenizer):
 # tokenize the text
 inputs = tokenizer(text, return_tensors="pt", max_length=512, truncation=True)
 # generate the embedding
 with torch.no_grad():
 outputs = model.encoder(input_ids=inputs.input_ids, attention_mask=inputs.attention_mask)
 # get the last hidden state
 last_hidden_states = outputs.encoder_last_hidden_state
 
 return last_hidden_states
 return embeddings

torch.cuda.is_available() returns true

I want to calculate the embeddings on the GPU using CUDA.

Answer 1

You need to explicitly move the model and the model inputs to the GPU.

You can run nvidia-smi to verify things are running on the GPU.

import torch
from transformers import BartTokenizer, BartForConditionalGeneration
model_name = 'facebook/bart-base'
device = 'cuda'
def load_model():
 # load the pretrained bart model
 model = BartForConditionalGeneration.from_pretrained(model_name)
 model.to(device)
 # load the tokenizer
 tokenizer = BartTokenizer.from_pretrained(model_name)
 return model, tokenizer
def calculate_text_embeddings(text, model, tokenizer):
 # tokenize the text
 inputs = tokenizer(text, return_tensors="pt", max_length=512, truncation=True)
 # generate the embedding
 with torch.no_grad():
 outputs = model.encoder(input_ids=inputs.input_ids.to(device), attention_mask=inputs.attention_mask.to(device))
 # get the last hidden state
 last_hidden_states = outputs.encoder_last_hidden_state
 
 return last_hidden_states.cpu()