Update README.md

README.md

@@ -25,7 +25,7 @@ This model leverages IndoBERT for understanding language and a Long Short-Term M
 ```python
 import torch.nn as nn
 from huggingface_hub import PyTorchModelHubMixin
-from transformers import AutoModelForSequenceClassification, AutoTokenizer
+from transformers import BertModel, AutoTokenizer
 
 class IndoBertLSTMEcommerceReview(nn.Module, PyTorchModelHubMixin):
     def __init__(self, bert):
@@ -44,9 +44,7 @@ class IndoBertLSTMEcommerceReview(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-lstm-finetuned-ecommerce-reviews")
 model = IndoBertLSTMEcommerceReview.from_pretrained("fahrendrakhoirul/indobert-lstm-finetuned-ecommerce-reviews", bert=bert)
   ```