Update README.md

README.md

@@ -26,7 +26,7 @@ Clearly define the three output classes and their corresponding labels:
 ```python
 import torch.nn as nn
 from huggingface_hub import PyTorchModelHubMixin
-from transformers import AutoModelForSequenceClassification, AutoTokenizer
+from transformers import BertModel, AutoTokenizer
 
 class IndoBertCNNEcommerceReview(nn.Module, PyTorchModelHubMixin):
         def __init__(self, bert):
@@ -50,7 +50,7 @@ class IndoBertCNNEcommerceReview(nn.Module, PyTorchModelHubMixin):
             probabilities = self.sigmoid(logits)
             return probabilities
 
-bert = AutoModelForSequenceClassification.from_pretrained("indobenchmark/indobert-base-p1",
+bert = BertModel.from_pretrained("indobenchmark/indobert-base-p1",
                                                             num_labels=3,
                                                            problem_type="multi_label_classification")
 tokenizer = AutoTokenizer.from_pretrained("fahrendrakhoirul/indobert-cnn-finetuned-ecommerce-reviews")