Update README.md

README.md

@@ -6,7 +6,7 @@ inference: false
 A simple use case:
 
 ```shell
-from transformers import Wav2Vec2Processor, AutoModel
+from transformers import Wav2Vec2Processor, HubertModel
 import torch
 from torch import nn
 from datasets import load_dataset
@@ -15,10 +15,10 @@ from datasets import load_dataset
 dataset = load_dataset("hf-internal-testing/librispeech_asr_demo", "clean", split="validation")
 dataset = dataset.sort("id")
 sampling_rate = dataset.features["audio"].sampling_rate
-processor = Wav2Vec2Processor.from_pretrained("facebook/data2vec-audio-base-960h")
+processor = Wav2Vec2Processor.from_pretrained("facebook/hubert-large-ls960-ft")
 
 # loading our model weights
-model = AutoModel.from_pretrained("m-a-p/MERT-v0")
+model = HubertModel.from_pretrained("m-a-p/MERT-v0")
 
 # audio file is decoded on the fly
 inputs = processor(dataset[0]["audio"]["array"], sampling_rate=sampling_rate, return_tensors="pt")
@@ -36,6 +36,6 @@ print(time_reduced_hidden_states.shape) # [13, 768]
 
 # you can even use a learnable weighted average representation
 aggregator = nn.Conv1d(in_channels=13, out_channels=1, kernel_size=1)
-weighted_avg_hidden_states = aggregator(time_reduced_hidden_states).squeeze()
+weighted_avg_hidden_states = aggregator(time_reduced_hidden_states.unsqueeze(0)).squeeze()
 print(weighted_avg_hidden_states.shape) # [768]
 ```