diff --git a/setup.py b/setup.py
index 82a7aa9..a28d2e2 100644
--- a/setup.py
+++ b/setup.py
@@ -32,8 +32,8 @@
         "ray",
         'numpy',
         'urlextract',
-        "transformers<=4.21.2",  # push-to-model is not working for latest version
-        "huggingface-hub<=0.9.1",
+        "transformers==4.37.2",
+        "huggingface-hub",
         'sentence_transformers',
         'torch',
         'datasets'
diff --git a/tweetnlp/text_classification/trainer.py b/tweetnlp/text_classification/trainer.py
index 65d2df3..e4d4b8c 100644
--- a/tweetnlp/text_classification/trainer.py
+++ b/tweetnlp/text_classification/trainer.py
@@ -10,7 +10,7 @@
 import torch
 import numpy as np
 from huggingface_hub import create_repo
-from datasets import load_metric
+from datasets import load_dataset
 from datasets.dataset_dict import DatasetDict
 from transformers import TrainingArguments, Trainer
 from ray import tune