Update README.md

README.md

@@ -55,15 +55,14 @@ Further details are available in the corresponding [**paper**](https://arxiv.org
 ```python 
 import torch
 import torch.nn as nn
-from transformers import HubertForSequenceClassification, Wav2Vec2FeatureExtractor
+from transformers import AutoModelForAudioClassification, Wav2Vec2FeatureExtractor
 
 
 
 # CONFIG and MODEL SETUP
 model_name = 'amiriparian/HuBERT-EmoSet'
 feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained("facebook/hubert-base-ls960")
-model = HubertForSequenceClassification.from_pretrained(model_name)
-model.classifier = nn.Linear(in_features=256,out_features=6)
+AutoModelForAudioClassification.from_pretrained(model_name, trust_remote_code=True)
 
 sampling_rate=16000 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")