ayushshah
/

distilbert-dapt-imdb-sentiment

Text Classification

sentiment-analysis

domain-adaptation

Eval Results (legacy)

Model card Files Files and versions

ayushshah commited on Jan 4

Commit

12a6660

·

1 Parent(s): c9669c2

Upload model.py

Files changed (1) hide show

model.py +1 -1

model.py CHANGED Viewed

@@ -23,7 +23,7 @@ def infer_reviews(reviews, model, tokenizer):
     if is_single:
         reviews = [reviews]
-    inputs = tokenizer(reviews, return_tensors="pt", truncation=True, padding=True, max_length=512).to('cuda')
     with torch.no_grad():
         logits = model(inputs['input_ids'], inputs['attention_mask'])

     if is_single:
         reviews = [reviews]
+    inputs = tokenizer(reviews, return_tensors="pt", truncation=True, padding=True, max_length=512).to('cuda' if torch.cuda.is_available() else 'cpu')
     with torch.no_grad():
         logits = model(inputs['input_ids'], inputs['attention_mask'])