Transfer Learning using HuggingFace and Tensorflow with AutoModel does not work

Question

I try to do transfer learning using a HuggingFace pretrained BERT model. I want to use tensorflow API with it. I do not understand why the last line produces an error

from transformers import AutoTokenizer, AutoModel

model_name = "distilbert-base-uncased"
text = "this is a test"
tokenizer = AutoTokenizer.from_pretrained(model_name)    
text_tensor = tokenizer.encode(text, return_tensors="tf")

model = AutoModel.from_pretrained(model_name).to("cuda")
output = model(text_tensor) # ERROR!!, but why?

Answer 1

You are mixing Tensorflow and Pytorch.

Use TFAutoModel instead of default (Pytorch) AutoModel

from transformers import AutoTokenizer, TFAutoModel

model_name = "distilbert-base-uncased"
text = "this is a test"
tokenizer = AutoTokenizer.from_pretrained(model_name)    
text_tensor = tokenizer.encode(text, return_tensors="tf")

model = TFAutoModel.from_pretrained(model_name).to("cuda")
output = model(text_tensor)