Update README.md

README.md

@@ -92,7 +92,7 @@ class AgeGenderModel(Wav2Vec2PreTrainedModel):
         hidden_states = outputs[0]
         hidden_states = torch.mean(hidden_states, dim=1)
         logits_age = self.age(hidden_states)
-        logits_gender = self.gender(hidden_states)
+        logits_gender = torch.softmax(self.gender(hidden_states), dim=1)
 
         return hidden_states, logits_age, logits_gender
 
@@ -140,10 +140,10 @@ def process_func(
 
 print(process_func(signal, sampling_rate))
 #    Age       child      female      male
-# [[ 0.3079211 -1.6096017 -2.1094327  3.1461434]]
+# [[ 0.3079211  0.00848487 0.0051472  0.9863679 ]]
 
 print(process_func(signal, sampling_rate, embeddings=True))
 # Pooled hidden states of last transformer layer
-# [[-0.00752167  0.0065819  -0.00746342 ...  0.00663632  0.00848748
-#    0.00599211]]
+# [[ 0.00409924  0.00365688  0.02392936 ...  0.02349018 -0.13294911
+#    0.1538802 ]]
 ```