Update README.md

README.md

@@ -35,7 +35,7 @@ Optionally, provide brief examples of reviews that would fall into each category
 ```python
 import torch.nn as nn
 from huggingface_hub import PyTorchModelHubMixin
-from transformers import AutoModelForSequenceClassification, AutoTokenizer
+from transformers import BertModel, AutoTokenizer
 
 class IndoBertEcommerceReview(nn.Module, PyTorchModelHubMixin):
         def __init__(self, bert):
@@ -49,9 +49,7 @@ class IndoBertEcommerceReview(nn.Module, PyTorchModelHubMixin):
             probabilities = self.sigmoid(logits)
             return probabilities
 
-bert = AutoModelForSequenceClassification.from_pretrained("indobenchmark/indobert-base-p1",
-                                                            num_labels=3,
-                                                           problem_type="multi_label_classification")
+bert = BertModel.from_pretrained("indobenchmark/indobert-base-p1")
 tokenizer = AutoTokenizer.from_pretrained("fahrendrakhoirul/indobert-finetuned-ecommerce-reviews")
 model = IndoBertEcommerceReview.from_pretrained("fahrendrakhoirul/indobert-finetuned-ecommerce-reviews", bert=bert)
   ```