Update README.md

README.md

@@ -2684,7 +2684,7 @@ def pooling(outputs: torch.Tensor, inputs: Dict,  strategy: str = 'cls') -> np.n
         outputs = outputs[:, 0]
     elif strategy == 'mean':
         outputs = torch.sum(
-            outputs * inputs["attention_mask"][:, :, None], dim=1) / torch.sum(inputs["attention_mask"])
+            outputs * inputs["attention_mask"][:, :, None], dim=1) / torch.sum(inputs["attention_mask"], dim=1, keepdim=True)
     else:
         raise NotImplementedError
     return outputs.detach().cpu().numpy()