vkovacs commited on
Commit
f61fe0b
·
1 Parent(s): c1f3033

device map auto

Browse files
Files changed (1) hide show
  1. app.py +1 -3
app.py CHANGED
@@ -41,16 +41,14 @@ def build_huggingface_path(language: str):
41
  return "poltextlab/xlm-roberta-large-pooled-MORES"
42
 
43
  def predict(text, model_id, tokenizer_id):
44
- device = torch.device("cpu")
45
  model = AutoModelForSequenceClassification.from_pretrained(model_id, low_cpu_mem_usage=True, device_map="auto", offload_folder="offload", token=HF_TOKEN)
46
  tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
47
- model.to(device)
48
 
49
  inputs = tokenizer(text,
50
  max_length=64,
51
  truncation=True,
52
  padding="do_not_pad",
53
- return_tensors="pt").to(device)
54
  model.eval()
55
 
56
  with torch.no_grad():
 
41
  return "poltextlab/xlm-roberta-large-pooled-MORES"
42
 
43
  def predict(text, model_id, tokenizer_id):
 
44
  model = AutoModelForSequenceClassification.from_pretrained(model_id, low_cpu_mem_usage=True, device_map="auto", offload_folder="offload", token=HF_TOKEN)
45
  tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
 
46
 
47
  inputs = tokenizer(text,
48
  max_length=64,
49
  truncation=True,
50
  padding="do_not_pad",
51
+ return_tensors="pt")
52
  model.eval()
53
 
54
  with torch.no_grad():