We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
The comment in onmt.modules.average_attn may be wrong
The comment in onmt.modules.average_attn.AverageAttention.forward says that it gets a input of size `[batch_size x input_len x model_dim]`, and outputs a gating_outputs and a average_outputs both with size `[batch_size x 1 x model_dim]`. However, in my test the two outputs are both with size [batch_size x input_len x model_dim].
```python
import torch
import onmt.modules.average_attn as Att
model = Att.AverageAttention(3)
a = torch.arange(24).reshape(2,4,3).float()
print("intput size:", a.size())
o1, o2 = model(a)
print("o1 size:", o1.size())
print("o2.size:", o2.size())
```
intput size: torch.Size([2, 4, 3])
o1 size: torch.Size([2, 4, 3])
o2 size: torch.Size([2, 4, 3])

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `onmt/modules/average_attn.py`
Content:
```
1 # -*- coding: utf-8 -*-
2 """ Average Attention module """
3 
4 import torch
5 import torch.nn as nn
6 
7 from onmt.modules.position_ffn import PositionwiseFeedForward
8 
9 
10 class AverageAttention(nn.Module):
11     """
12     Average Attention module from
13     "Accelerating Neural Transformer via an Average Attention Network"
14     :cite:`https://arxiv.org/abs/1805.00631`.
15 
16     Args:
17        model_dim (int): the dimension of keys/values/queries,
18            must be divisible by head_count
19        dropout (float): dropout parameter
20     """
21 
22     def __init__(self, model_dim, dropout=0.1):
23         self.model_dim = model_dim
24 
25         super(AverageAttention, self).__init__()
26 
27         self.average_layer = PositionwiseFeedForward(model_dim, model_dim,
28                                                      dropout)
29         self.gating_layer = nn.Linear(model_dim * 2, model_dim * 2)
30 
31     def cumulative_average_mask(self, batch_size, inputs_len):
32         """
33         Builds the mask to compute the cumulative average as described in
34         https://arxiv.org/abs/1805.00631 -- Figure 3
35 
36         Args:
37             batch_size (int): batch size
38             inputs_len (int): length of the inputs
39 
40         Returns:
41             (`FloatTensor`):
42 
43             * A Tensor of shape `[batch_size x input_len x input_len]`
44         """
45 
46         triangle = torch.tril(torch.ones(inputs_len, inputs_len))
47         weights = torch.ones(1, inputs_len) / torch.arange(
48             1, inputs_len + 1, dtype=torch.float)
49         mask = triangle * weights.transpose(0, 1)
50 
51         return mask.unsqueeze(0).expand(batch_size, inputs_len, inputs_len)
52 
53     def cumulative_average(self, inputs, mask_or_step,
54                            layer_cache=None, step=None):
55         """
56         Computes the cumulative average as described in
57         https://arxiv.org/abs/1805.00631 -- Equations (1) (5) (6)
58 
59         Args:
60             inputs (`FloatTensor`): sequence to average
61                 `[batch_size x input_len x dimension]`
62             mask_or_step: if cache is set, this is assumed
63                 to be the current step of the
64                 dynamic decoding. Otherwise, it is the mask matrix
65                 used to compute the cumulative average.
66             cache: a dictionary containing the cumulative average
67                 of the previous step.
68         """
69         if layer_cache is not None:
70             step = mask_or_step
71             device = inputs.device
72             average_attention = (inputs + step *
73                                  layer_cache["prev_g"].to(device)) / (step + 1)
74             layer_cache["prev_g"] = average_attention
75             return average_attention
76         else:
77             mask = mask_or_step
78             return torch.matmul(mask, inputs)
79 
80     def forward(self, inputs, mask=None, layer_cache=None, step=None):
81         """
82         Args:
83             inputs (`FloatTensor`): `[batch_size x input_len x model_dim]`
84 
85         Returns:
86             (`FloatTensor`, `FloatTensor`):
87 
88             * gating_outputs `[batch_size x 1 x model_dim]`
89             * average_outputs average attention `[batch_size x 1 x model_dim]`
90         """
91         batch_size = inputs.size(0)
92         inputs_len = inputs.size(1)
93 
94         device = inputs.device
95         average_outputs = self.cumulative_average(
96           inputs, self.cumulative_average_mask(batch_size,
97                                                inputs_len).to(device).float()
98           if layer_cache is None else step, layer_cache=layer_cache)
99         average_outputs = self.average_layer(average_outputs)
100         gating_outputs = self.gating_layer(torch.cat((inputs,
101                                                       average_outputs), -1))
102         input_gate, forget_gate = torch.chunk(gating_outputs, 2, dim=2)
103         gating_outputs = torch.sigmoid(input_gate) * inputs + \
104             torch.sigmoid(forget_gate) * average_outputs
105 
106         return gating_outputs, average_outputs
107
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/onmt/modules/average_attn.py b/onmt/modules/average_attn.py
--- a/onmt/modules/average_attn.py
+++ b/onmt/modules/average_attn.py
@@ -65,6 +65,9 @@
                 used to compute the cumulative average.
             cache: a dictionary containing the cumulative average
                 of the previous step.
+
+        Returns:
+            a tensor of the same shape and type as :obj:`inputs`.
         """
         if layer_cache is not None:
             step = mask_or_step
@@ -85,8 +88,9 @@
         Returns:
             (`FloatTensor`, `FloatTensor`):
 
-            * gating_outputs `[batch_size x 1 x model_dim]`
-            * average_outputs average attention `[batch_size x 1 x model_dim]`
+            * gating_outputs `[batch_size x input_len x model_dim]`
+            * average_outputs average attention
+                `[batch_size x input_len x model_dim]`
         """
         batch_size = inputs.size(0)
         inputs_len = inputs.size(1)

{"golden_diff": "diff --git a/onmt/modules/average_attn.py b/onmt/modules/average_attn.py\n--- a/onmt/modules/average_attn.py\n+++ b/onmt/modules/average_attn.py\n@@ -65,6 +65,9 @@\n                 used to compute the cumulative average.\n             cache: a dictionary containing the cumulative average\n                 of the previous step.\n+\n+        Returns:\n+            a tensor of the same shape and type as :obj:`inputs`.\n         \"\"\"\n         if layer_cache is not None:\n             step = mask_or_step\n@@ -85,8 +88,9 @@\n         Returns:\n             (`FloatTensor`, `FloatTensor`):\n \n-            * gating_outputs `[batch_size x 1 x model_dim]`\n-            * average_outputs average attention `[batch_size x 1 x model_dim]`\n+            * gating_outputs `[batch_size x input_len x model_dim]`\n+            * average_outputs average attention\n+                `[batch_size x input_len x model_dim]`\n         \"\"\"\n         batch_size = inputs.size(0)\n         inputs_len = inputs.size(1)\n", "issue": "The comment in onmt.modules.average_attn may be wrong\nThe comment in onmt.modules.average_attn.AverageAttention.forward says that it gets a input of size `[batch_size x input_len x model_dim]`, and outputs a gating_outputs and a average_outputs both with size `[batch_size x 1 x model_dim]`. However, in my test the two outputs are both with size [batch_size x input_len x model_dim].\r\n```python\r\nimport torch\r\nimport onmt.modules.average_attn as Att\r\nmodel = Att.AverageAttention(3)\r\na = torch.arange(24).reshape(2,4,3).float()\r\nprint(\"intput size:\", a.size())\r\no1, o2 = model(a)\r\nprint(\"o1 size:\", o1.size())\r\nprint(\"o2.size:\", o2.size())\r\n```\r\nintput size: torch.Size([2, 4, 3])\r\no1 size: torch.Size([2, 4, 3])\r\no2 size: torch.Size([2, 4, 3])\n", "before_files": [{"content": "# -*- coding: utf-8 -*-\n\"\"\" Average Attention module \"\"\"\n\nimport torch\nimport torch.nn as nn\n\nfrom onmt.modules.position_ffn import PositionwiseFeedForward\n\n\nclass AverageAttention(nn.Module):\n    \"\"\"\n    Average Attention module from\n    \"Accelerating Neural Transformer via an Average Attention Network\"\n    :cite:`https://arxiv.org/abs/1805.00631`.\n\n    Args:\n       model_dim (int): the dimension of keys/values/queries,\n           must be divisible by head_count\n       dropout (float): dropout parameter\n    \"\"\"\n\n    def __init__(self, model_dim, dropout=0.1):\n        self.model_dim = model_dim\n\n        super(AverageAttention, self).__init__()\n\n        self.average_layer = PositionwiseFeedForward(model_dim, model_dim,\n                                                     dropout)\n        self.gating_layer = nn.Linear(model_dim * 2, model_dim * 2)\n\n    def cumulative_average_mask(self, batch_size, inputs_len):\n        \"\"\"\n        Builds the mask to compute the cumulative average as described in\n        https://arxiv.org/abs/1805.00631 -- Figure 3\n\n        Args:\n            batch_size (int): batch size\n            inputs_len (int): length of the inputs\n\n        Returns:\n            (`FloatTensor`):\n\n            * A Tensor of shape `[batch_size x input_len x input_len]`\n        \"\"\"\n\n        triangle = torch.tril(torch.ones(inputs_len, inputs_len))\n        weights = torch.ones(1, inputs_len) / torch.arange(\n            1, inputs_len + 1, dtype=torch.float)\n        mask = triangle * weights.transpose(0, 1)\n\n        return mask.unsqueeze(0).expand(batch_size, inputs_len, inputs_len)\n\n    def cumulative_average(self, inputs, mask_or_step,\n                           layer_cache=None, step=None):\n        \"\"\"\n        Computes the cumulative average as described in\n        https://arxiv.org/abs/1805.00631 -- Equations (1) (5) (6)\n\n        Args:\n            inputs (`FloatTensor`): sequence to average\n                `[batch_size x input_len x dimension]`\n            mask_or_step: if cache is set, this is assumed\n                to be the current step of the\n                dynamic decoding. Otherwise, it is the mask matrix\n                used to compute the cumulative average.\n            cache: a dictionary containing the cumulative average\n                of the previous step.\n        \"\"\"\n        if layer_cache is not None:\n            step = mask_or_step\n            device = inputs.device\n            average_attention = (inputs + step *\n                                 layer_cache[\"prev_g\"].to(device)) / (step + 1)\n            layer_cache[\"prev_g\"] = average_attention\n            return average_attention\n        else:\n            mask = mask_or_step\n            return torch.matmul(mask, inputs)\n\n    def forward(self, inputs, mask=None, layer_cache=None, step=None):\n        \"\"\"\n        Args:\n            inputs (`FloatTensor`): `[batch_size x input_len x model_dim]`\n\n        Returns:\n            (`FloatTensor`, `FloatTensor`):\n\n            * gating_outputs `[batch_size x 1 x model_dim]`\n            * average_outputs average attention `[batch_size x 1 x model_dim]`\n        \"\"\"\n        batch_size = inputs.size(0)\n        inputs_len = inputs.size(1)\n\n        device = inputs.device\n        average_outputs = self.cumulative_average(\n          inputs, self.cumulative_average_mask(batch_size,\n                                               inputs_len).to(device).float()\n          if layer_cache is None else step, layer_cache=layer_cache)\n        average_outputs = self.average_layer(average_outputs)\n        gating_outputs = self.gating_layer(torch.cat((inputs,\n                                                      average_outputs), -1))\n        input_gate, forget_gate = torch.chunk(gating_outputs, 2, dim=2)\n        gating_outputs = torch.sigmoid(input_gate) * inputs + \\\n            torch.sigmoid(forget_gate) * average_outputs\n\n        return gating_outputs, average_outputs\n", "path": "onmt/modules/average_attn.py"}], "after_files": [{"content": "# -*- coding: utf-8 -*-\n\"\"\" Average Attention module \"\"\"\n\nimport torch\nimport torch.nn as nn\n\nfrom onmt.modules.position_ffn import PositionwiseFeedForward\n\n\nclass AverageAttention(nn.Module):\n    \"\"\"\n    Average Attention module from\n    \"Accelerating Neural Transformer via an Average Attention Network\"\n    :cite:`https://arxiv.org/abs/1805.00631`.\n\n    Args:\n       model_dim (int): the dimension of keys/values/queries,\n           must be divisible by head_count\n       dropout (float): dropout parameter\n    \"\"\"\n\n    def __init__(self, model_dim, dropout=0.1):\n        self.model_dim = model_dim\n\n        super(AverageAttention, self).__init__()\n\n        self.average_layer = PositionwiseFeedForward(model_dim, model_dim,\n                                                     dropout)\n        self.gating_layer = nn.Linear(model_dim * 2, model_dim * 2)\n\n    def cumulative_average_mask(self, batch_size, inputs_len):\n        \"\"\"\n        Builds the mask to compute the cumulative average as described in\n        https://arxiv.org/abs/1805.00631 -- Figure 3\n\n        Args:\n            batch_size (int): batch size\n            inputs_len (int): length of the inputs\n\n        Returns:\n            (`FloatTensor`):\n\n            * A Tensor of shape `[batch_size x input_len x input_len]`\n        \"\"\"\n\n        triangle = torch.tril(torch.ones(inputs_len, inputs_len))\n        weights = torch.ones(1, inputs_len) / torch.arange(\n            1, inputs_len + 1, dtype=torch.float)\n        mask = triangle * weights.transpose(0, 1)\n\n        return mask.unsqueeze(0).expand(batch_size, inputs_len, inputs_len)\n\n    def cumulative_average(self, inputs, mask_or_step,\n                           layer_cache=None, step=None):\n        \"\"\"\n        Computes the cumulative average as described in\n        https://arxiv.org/abs/1805.00631 -- Equations (1) (5) (6)\n\n        Args:\n            inputs (`FloatTensor`): sequence to average\n                `[batch_size x input_len x dimension]`\n            mask_or_step: if cache is set, this is assumed\n                to be the current step of the\n                dynamic decoding. Otherwise, it is the mask matrix\n                used to compute the cumulative average.\n            cache: a dictionary containing the cumulative average\n                of the previous step.\n\n        Returns:\n            a tensor of the same shape and type as :obj:`inputs`.\n        \"\"\"\n        if layer_cache is not None:\n            step = mask_or_step\n            device = inputs.device\n            average_attention = (inputs + step *\n                                 layer_cache[\"prev_g\"].to(device)) / (step + 1)\n            layer_cache[\"prev_g\"] = average_attention\n            return average_attention\n        else:\n            mask = mask_or_step\n            return torch.matmul(mask, inputs)\n\n    def forward(self, inputs, mask=None, layer_cache=None, step=None):\n        \"\"\"\n        Args:\n            inputs (`FloatTensor`): `[batch_size x input_len x model_dim]`\n\n        Returns:\n            (`FloatTensor`, `FloatTensor`):\n\n            * gating_outputs `[batch_size x input_len x model_dim]`\n            * average_outputs average attention\n                `[batch_size x input_len x model_dim]`\n        \"\"\"\n        batch_size = inputs.size(0)\n        inputs_len = inputs.size(1)\n\n        device = inputs.device\n        average_outputs = self.cumulative_average(\n          inputs, self.cumulative_average_mask(batch_size,\n                                               inputs_len).to(device).float()\n          if layer_cache is None else step, layer_cache=layer_cache)\n        average_outputs = self.average_layer(average_outputs)\n        gating_outputs = self.gating_layer(torch.cat((inputs,\n                                                      average_outputs), -1))\n        input_gate, forget_gate = torch.chunk(gating_outputs, 2, dim=2)\n        gating_outputs = torch.sigmoid(input_gate) * inputs + \\\n            torch.sigmoid(forget_gate) * average_outputs\n\n        return gating_outputs, average_outputs\n", "path": "onmt/modules/average_attn.py"}]}