We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Repeated check in code
https://github.com/OpenNMT/OpenNMT-py/blob/master/onmt/modules/multi_headed_attn.py#L133, why is the layer_cache is not None check required again? Similarly will this be ever called? https://github.com/OpenNMT/OpenNMT-py/blob/master/onmt/modules/multi_headed_attn.py#L158 

--- 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/multi_headed_attn.py`
Content:
```
1 """ Multi-Head Attention module """
2 import math
3 import torch
4 import torch.nn as nn
5 
6 # from onmt.utils.misc import aeq
7 
8 
9 class MultiHeadedAttention(nn.Module):
10     """
11     Multi-Head Attention module from
12     "Attention is All You Need"
13     :cite:`DBLP:journals/corr/VaswaniSPUJGKP17`.
14 
15     Similar to standard `dot` attention but uses
16     multiple attention distributions simulataneously
17     to select relevant items.
18 
19     .. mermaid::
20 
21        graph BT
22           A[key]
23           B[value]
24           C[query]
25           O[output]
26           subgraph Attn
27             D[Attn 1]
28             E[Attn 2]
29             F[Attn N]
30           end
31           A --> D
32           C --> D
33           A --> E
34           C --> E
35           A --> F
36           C --> F
37           D --> O
38           E --> O
39           F --> O
40           B --> O
41 
42     Also includes several additional tricks.
43 
44     Args:
45        head_count (int): number of parallel heads
46        model_dim (int): the dimension of keys/values/queries,
47            must be divisible by head_count
48        dropout (float): dropout parameter
49     """
50 
51     def __init__(self, head_count, model_dim, dropout=0.1):
52         assert model_dim % head_count == 0
53         self.dim_per_head = model_dim // head_count
54         self.model_dim = model_dim
55 
56         super(MultiHeadedAttention, self).__init__()
57         self.head_count = head_count
58 
59         self.linear_keys = nn.Linear(model_dim,
60                                      head_count * self.dim_per_head)
61         self.linear_values = nn.Linear(model_dim,
62                                        head_count * self.dim_per_head)
63         self.linear_query = nn.Linear(model_dim,
64                                       head_count * self.dim_per_head)
65         self.softmax = nn.Softmax(dim=-1)
66         self.dropout = nn.Dropout(dropout)
67         self.final_linear = nn.Linear(model_dim, model_dim)
68 
69     def forward(self, key, value, query, mask=None,
70                 layer_cache=None, type=None):
71         """
72         Compute the context vector and the attention vectors.
73 
74         Args:
75            key (`FloatTensor`): set of `key_len`
76                 key vectors `[batch, key_len, dim]`
77            value (`FloatTensor`): set of `key_len`
78                 value vectors `[batch, key_len, dim]`
79            query (`FloatTensor`): set of `query_len`
80                  query vectors  `[batch, query_len, dim]`
81            mask: binary mask indicating which keys have
82                  non-zero attention `[batch, query_len, key_len]`
83         Returns:
84            (`FloatTensor`, `FloatTensor`) :
85 
86            * output context vectors `[batch, query_len, dim]`
87            * one of the attention vectors `[batch, query_len, key_len]`
88         """
89 
90         # CHECKS
91         # batch, k_len, d = key.size()
92         # batch_, k_len_, d_ = value.size()
93         # aeq(batch, batch_)
94         # aeq(k_len, k_len_)
95         # aeq(d, d_)
96         # batch_, q_len, d_ = query.size()
97         # aeq(batch, batch_)
98         # aeq(d, d_)
99         # aeq(self.model_dim % 8, 0)
100         # if mask is not None:
101         #    batch_, q_len_, k_len_ = mask.size()
102         #    aeq(batch_, batch)
103         #    aeq(k_len_, k_len)
104         #    aeq(q_len_ == q_len)
105         # END CHECKS
106 
107         batch_size = key.size(0)
108         dim_per_head = self.dim_per_head
109         head_count = self.head_count
110         key_len = key.size(1)
111         query_len = query.size(1)
112 
113         def shape(x):
114             """  projection """
115             return x.view(batch_size, -1, head_count, dim_per_head) \
116                 .transpose(1, 2)
117 
118         def unshape(x):
119             """  compute context """
120             return x.transpose(1, 2).contiguous() \
121                     .view(batch_size, -1, head_count * dim_per_head)
122 
123         # 1) Project key, value, and query.
124         if layer_cache is not None:
125             if type == "self":
126                 query, key, value = self.linear_query(query),\
127                                     self.linear_keys(query),\
128                                     self.linear_values(query)
129 
130                 key = shape(key)
131                 value = shape(value)
132 
133                 if layer_cache is not None:
134                     device = key.device
135                     if layer_cache["self_keys"] is not None:
136                         key = torch.cat(
137                             (layer_cache["self_keys"].to(device), key),
138                             dim=2)
139                     if layer_cache["self_values"] is not None:
140                         value = torch.cat(
141                             (layer_cache["self_values"].to(device), value),
142                             dim=2)
143                     layer_cache["self_keys"] = key
144                     layer_cache["self_values"] = value
145             elif type == "context":
146                 query = self.linear_query(query)
147                 if layer_cache is not None:
148                     if layer_cache["memory_keys"] is None:
149                         key, value = self.linear_keys(key),\
150                                      self.linear_values(value)
151                         key = shape(key)
152                         value = shape(value)
153                     else:
154                         key, value = layer_cache["memory_keys"],\
155                                    layer_cache["memory_values"]
156                     layer_cache["memory_keys"] = key
157                     layer_cache["memory_values"] = value
158                 else:
159                     key, value = self.linear_keys(key),\
160                                  self.linear_values(value)
161                     key = shape(key)
162                     value = shape(value)
163         else:
164             key = self.linear_keys(key)
165             value = self.linear_values(value)
166             query = self.linear_query(query)
167             key = shape(key)
168             value = shape(value)
169 
170         query = shape(query)
171 
172         key_len = key.size(2)
173         query_len = query.size(2)
174 
175         # 2) Calculate and scale scores.
176         query = query / math.sqrt(dim_per_head)
177         scores = torch.matmul(query, key.transpose(2, 3))
178 
179         if mask is not None:
180             mask = mask.unsqueeze(1)  # [B, 1, 1, T_values]
181             scores = scores.masked_fill(mask, -1e18)
182 
183         # 3) Apply attention dropout and compute context vectors.
184         attn = self.softmax(scores)
185         drop_attn = self.dropout(attn)
186         context = unshape(torch.matmul(drop_attn, value))
187 
188         output = self.final_linear(context)
189         # CHECK
190         # batch_, q_len_, d_ = output.size()
191         # aeq(q_len, q_len_)
192         # aeq(batch, batch_)
193         # aeq(d, d_)
194 
195         # Return one attn
196         top_attn = attn \
197             .view(batch_size, head_count,
198                   query_len, key_len)[:, 0, :, :] \
199             .contiguous()
200 
201         return output, top_attn
202
```
--- 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/multi_headed_attn.py b/onmt/modules/multi_headed_attn.py
--- a/onmt/modules/multi_headed_attn.py
+++ b/onmt/modules/multi_headed_attn.py
@@ -126,40 +126,31 @@
                 query, key, value = self.linear_query(query),\
                                     self.linear_keys(query),\
                                     self.linear_values(query)
-
                 key = shape(key)
                 value = shape(value)
-
-                if layer_cache is not None:
-                    device = key.device
-                    if layer_cache["self_keys"] is not None:
-                        key = torch.cat(
-                            (layer_cache["self_keys"].to(device), key),
-                            dim=2)
-                    if layer_cache["self_values"] is not None:
-                        value = torch.cat(
-                            (layer_cache["self_values"].to(device), value),
-                            dim=2)
-                    layer_cache["self_keys"] = key
-                    layer_cache["self_values"] = value
+                device = key.device
+                if layer_cache["self_keys"] is not None:
+                    key = torch.cat(
+                        (layer_cache["self_keys"].to(device), key),
+                        dim=2)
+                if layer_cache["self_values"] is not None:
+                    value = torch.cat(
+                        (layer_cache["self_values"].to(device), value),
+                        dim=2)
+                layer_cache["self_keys"] = key
+                layer_cache["self_values"] = value
             elif type == "context":
                 query = self.linear_query(query)
-                if layer_cache is not None:
-                    if layer_cache["memory_keys"] is None:
-                        key, value = self.linear_keys(key),\
-                                     self.linear_values(value)
-                        key = shape(key)
-                        value = shape(value)
-                    else:
-                        key, value = layer_cache["memory_keys"],\
-                                   layer_cache["memory_values"]
-                    layer_cache["memory_keys"] = key
-                    layer_cache["memory_values"] = value
-                else:
+                if layer_cache["memory_keys"] is None:
                     key, value = self.linear_keys(key),\
                                  self.linear_values(value)
                     key = shape(key)
                     value = shape(value)
+                else:
+                    key, value = layer_cache["memory_keys"],\
+                               layer_cache["memory_values"]
+                layer_cache["memory_keys"] = key
+                layer_cache["memory_values"] = value
         else:
             key = self.linear_keys(key)
             value = self.linear_values(value)

{"golden_diff": "diff --git a/onmt/modules/multi_headed_attn.py b/onmt/modules/multi_headed_attn.py\n--- a/onmt/modules/multi_headed_attn.py\n+++ b/onmt/modules/multi_headed_attn.py\n@@ -126,40 +126,31 @@\n                 query, key, value = self.linear_query(query),\\\n                                     self.linear_keys(query),\\\n                                     self.linear_values(query)\n-\n                 key = shape(key)\n                 value = shape(value)\n-\n-                if layer_cache is not None:\n-                    device = key.device\n-                    if layer_cache[\"self_keys\"] is not None:\n-                        key = torch.cat(\n-                            (layer_cache[\"self_keys\"].to(device), key),\n-                            dim=2)\n-                    if layer_cache[\"self_values\"] is not None:\n-                        value = torch.cat(\n-                            (layer_cache[\"self_values\"].to(device), value),\n-                            dim=2)\n-                    layer_cache[\"self_keys\"] = key\n-                    layer_cache[\"self_values\"] = value\n+                device = key.device\n+                if layer_cache[\"self_keys\"] is not None:\n+                    key = torch.cat(\n+                        (layer_cache[\"self_keys\"].to(device), key),\n+                        dim=2)\n+                if layer_cache[\"self_values\"] is not None:\n+                    value = torch.cat(\n+                        (layer_cache[\"self_values\"].to(device), value),\n+                        dim=2)\n+                layer_cache[\"self_keys\"] = key\n+                layer_cache[\"self_values\"] = value\n             elif type == \"context\":\n                 query = self.linear_query(query)\n-                if layer_cache is not None:\n-                    if layer_cache[\"memory_keys\"] is None:\n-                        key, value = self.linear_keys(key),\\\n-                                     self.linear_values(value)\n-                        key = shape(key)\n-                        value = shape(value)\n-                    else:\n-                        key, value = layer_cache[\"memory_keys\"],\\\n-                                   layer_cache[\"memory_values\"]\n-                    layer_cache[\"memory_keys\"] = key\n-                    layer_cache[\"memory_values\"] = value\n-                else:\n+                if layer_cache[\"memory_keys\"] is None:\n                     key, value = self.linear_keys(key),\\\n                                  self.linear_values(value)\n                     key = shape(key)\n                     value = shape(value)\n+                else:\n+                    key, value = layer_cache[\"memory_keys\"],\\\n+                               layer_cache[\"memory_values\"]\n+                layer_cache[\"memory_keys\"] = key\n+                layer_cache[\"memory_values\"] = value\n         else:\n             key = self.linear_keys(key)\n             value = self.linear_values(value)\n", "issue": "Repeated check in code\nhttps://github.com/OpenNMT/OpenNMT-py/blob/master/onmt/modules/multi_headed_attn.py#L133, why is the layer_cache is not None check required again? Similarly will this be ever called? https://github.com/OpenNMT/OpenNMT-py/blob/master/onmt/modules/multi_headed_attn.py#L158 \n", "before_files": [{"content": "\"\"\" Multi-Head Attention module \"\"\"\nimport math\nimport torch\nimport torch.nn as nn\n\n# from onmt.utils.misc import aeq\n\n\nclass MultiHeadedAttention(nn.Module):\n    \"\"\"\n    Multi-Head Attention module from\n    \"Attention is All You Need\"\n    :cite:`DBLP:journals/corr/VaswaniSPUJGKP17`.\n\n    Similar to standard `dot` attention but uses\n    multiple attention distributions simulataneously\n    to select relevant items.\n\n    .. mermaid::\n\n       graph BT\n          A[key]\n          B[value]\n          C[query]\n          O[output]\n          subgraph Attn\n            D[Attn 1]\n            E[Attn 2]\n            F[Attn N]\n          end\n          A --> D\n          C --> D\n          A --> E\n          C --> E\n          A --> F\n          C --> F\n          D --> O\n          E --> O\n          F --> O\n          B --> O\n\n    Also includes several additional tricks.\n\n    Args:\n       head_count (int): number of parallel heads\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, head_count, model_dim, dropout=0.1):\n        assert model_dim % head_count == 0\n        self.dim_per_head = model_dim // head_count\n        self.model_dim = model_dim\n\n        super(MultiHeadedAttention, self).__init__()\n        self.head_count = head_count\n\n        self.linear_keys = nn.Linear(model_dim,\n                                     head_count * self.dim_per_head)\n        self.linear_values = nn.Linear(model_dim,\n                                       head_count * self.dim_per_head)\n        self.linear_query = nn.Linear(model_dim,\n                                      head_count * self.dim_per_head)\n        self.softmax = nn.Softmax(dim=-1)\n        self.dropout = nn.Dropout(dropout)\n        self.final_linear = nn.Linear(model_dim, model_dim)\n\n    def forward(self, key, value, query, mask=None,\n                layer_cache=None, type=None):\n        \"\"\"\n        Compute the context vector and the attention vectors.\n\n        Args:\n           key (`FloatTensor`): set of `key_len`\n                key vectors `[batch, key_len, dim]`\n           value (`FloatTensor`): set of `key_len`\n                value vectors `[batch, key_len, dim]`\n           query (`FloatTensor`): set of `query_len`\n                 query vectors  `[batch, query_len, dim]`\n           mask: binary mask indicating which keys have\n                 non-zero attention `[batch, query_len, key_len]`\n        Returns:\n           (`FloatTensor`, `FloatTensor`) :\n\n           * output context vectors `[batch, query_len, dim]`\n           * one of the attention vectors `[batch, query_len, key_len]`\n        \"\"\"\n\n        # CHECKS\n        # batch, k_len, d = key.size()\n        # batch_, k_len_, d_ = value.size()\n        # aeq(batch, batch_)\n        # aeq(k_len, k_len_)\n        # aeq(d, d_)\n        # batch_, q_len, d_ = query.size()\n        # aeq(batch, batch_)\n        # aeq(d, d_)\n        # aeq(self.model_dim % 8, 0)\n        # if mask is not None:\n        #    batch_, q_len_, k_len_ = mask.size()\n        #    aeq(batch_, batch)\n        #    aeq(k_len_, k_len)\n        #    aeq(q_len_ == q_len)\n        # END CHECKS\n\n        batch_size = key.size(0)\n        dim_per_head = self.dim_per_head\n        head_count = self.head_count\n        key_len = key.size(1)\n        query_len = query.size(1)\n\n        def shape(x):\n            \"\"\"  projection \"\"\"\n            return x.view(batch_size, -1, head_count, dim_per_head) \\\n                .transpose(1, 2)\n\n        def unshape(x):\n            \"\"\"  compute context \"\"\"\n            return x.transpose(1, 2).contiguous() \\\n                    .view(batch_size, -1, head_count * dim_per_head)\n\n        # 1) Project key, value, and query.\n        if layer_cache is not None:\n            if type == \"self\":\n                query, key, value = self.linear_query(query),\\\n                                    self.linear_keys(query),\\\n                                    self.linear_values(query)\n\n                key = shape(key)\n                value = shape(value)\n\n                if layer_cache is not None:\n                    device = key.device\n                    if layer_cache[\"self_keys\"] is not None:\n                        key = torch.cat(\n                            (layer_cache[\"self_keys\"].to(device), key),\n                            dim=2)\n                    if layer_cache[\"self_values\"] is not None:\n                        value = torch.cat(\n                            (layer_cache[\"self_values\"].to(device), value),\n                            dim=2)\n                    layer_cache[\"self_keys\"] = key\n                    layer_cache[\"self_values\"] = value\n            elif type == \"context\":\n                query = self.linear_query(query)\n                if layer_cache is not None:\n                    if layer_cache[\"memory_keys\"] is None:\n                        key, value = self.linear_keys(key),\\\n                                     self.linear_values(value)\n                        key = shape(key)\n                        value = shape(value)\n                    else:\n                        key, value = layer_cache[\"memory_keys\"],\\\n                                   layer_cache[\"memory_values\"]\n                    layer_cache[\"memory_keys\"] = key\n                    layer_cache[\"memory_values\"] = value\n                else:\n                    key, value = self.linear_keys(key),\\\n                                 self.linear_values(value)\n                    key = shape(key)\n                    value = shape(value)\n        else:\n            key = self.linear_keys(key)\n            value = self.linear_values(value)\n            query = self.linear_query(query)\n            key = shape(key)\n            value = shape(value)\n\n        query = shape(query)\n\n        key_len = key.size(2)\n        query_len = query.size(2)\n\n        # 2) Calculate and scale scores.\n        query = query / math.sqrt(dim_per_head)\n        scores = torch.matmul(query, key.transpose(2, 3))\n\n        if mask is not None:\n            mask = mask.unsqueeze(1)  # [B, 1, 1, T_values]\n            scores = scores.masked_fill(mask, -1e18)\n\n        # 3) Apply attention dropout and compute context vectors.\n        attn = self.softmax(scores)\n        drop_attn = self.dropout(attn)\n        context = unshape(torch.matmul(drop_attn, value))\n\n        output = self.final_linear(context)\n        # CHECK\n        # batch_, q_len_, d_ = output.size()\n        # aeq(q_len, q_len_)\n        # aeq(batch, batch_)\n        # aeq(d, d_)\n\n        # Return one attn\n        top_attn = attn \\\n            .view(batch_size, head_count,\n                  query_len, key_len)[:, 0, :, :] \\\n            .contiguous()\n\n        return output, top_attn\n", "path": "onmt/modules/multi_headed_attn.py"}], "after_files": [{"content": "\"\"\" Multi-Head Attention module \"\"\"\nimport math\nimport torch\nimport torch.nn as nn\n\n# from onmt.utils.misc import aeq\n\n\nclass MultiHeadedAttention(nn.Module):\n    \"\"\"\n    Multi-Head Attention module from\n    \"Attention is All You Need\"\n    :cite:`DBLP:journals/corr/VaswaniSPUJGKP17`.\n\n    Similar to standard `dot` attention but uses\n    multiple attention distributions simulataneously\n    to select relevant items.\n\n    .. mermaid::\n\n       graph BT\n          A[key]\n          B[value]\n          C[query]\n          O[output]\n          subgraph Attn\n            D[Attn 1]\n            E[Attn 2]\n            F[Attn N]\n          end\n          A --> D\n          C --> D\n          A --> E\n          C --> E\n          A --> F\n          C --> F\n          D --> O\n          E --> O\n          F --> O\n          B --> O\n\n    Also includes several additional tricks.\n\n    Args:\n       head_count (int): number of parallel heads\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, head_count, model_dim, dropout=0.1):\n        assert model_dim % head_count == 0\n        self.dim_per_head = model_dim // head_count\n        self.model_dim = model_dim\n\n        super(MultiHeadedAttention, self).__init__()\n        self.head_count = head_count\n\n        self.linear_keys = nn.Linear(model_dim,\n                                     head_count * self.dim_per_head)\n        self.linear_values = nn.Linear(model_dim,\n                                       head_count * self.dim_per_head)\n        self.linear_query = nn.Linear(model_dim,\n                                      head_count * self.dim_per_head)\n        self.softmax = nn.Softmax(dim=-1)\n        self.dropout = nn.Dropout(dropout)\n        self.final_linear = nn.Linear(model_dim, model_dim)\n\n    def forward(self, key, value, query, mask=None,\n                layer_cache=None, type=None):\n        \"\"\"\n        Compute the context vector and the attention vectors.\n\n        Args:\n           key (`FloatTensor`): set of `key_len`\n                key vectors `[batch, key_len, dim]`\n           value (`FloatTensor`): set of `key_len`\n                value vectors `[batch, key_len, dim]`\n           query (`FloatTensor`): set of `query_len`\n                 query vectors  `[batch, query_len, dim]`\n           mask: binary mask indicating which keys have\n                 non-zero attention `[batch, query_len, key_len]`\n        Returns:\n           (`FloatTensor`, `FloatTensor`) :\n\n           * output context vectors `[batch, query_len, dim]`\n           * one of the attention vectors `[batch, query_len, key_len]`\n        \"\"\"\n\n        # CHECKS\n        # batch, k_len, d = key.size()\n        # batch_, k_len_, d_ = value.size()\n        # aeq(batch, batch_)\n        # aeq(k_len, k_len_)\n        # aeq(d, d_)\n        # batch_, q_len, d_ = query.size()\n        # aeq(batch, batch_)\n        # aeq(d, d_)\n        # aeq(self.model_dim % 8, 0)\n        # if mask is not None:\n        #    batch_, q_len_, k_len_ = mask.size()\n        #    aeq(batch_, batch)\n        #    aeq(k_len_, k_len)\n        #    aeq(q_len_ == q_len)\n        # END CHECKS\n\n        batch_size = key.size(0)\n        dim_per_head = self.dim_per_head\n        head_count = self.head_count\n        key_len = key.size(1)\n        query_len = query.size(1)\n\n        def shape(x):\n            \"\"\"  projection \"\"\"\n            return x.view(batch_size, -1, head_count, dim_per_head) \\\n                .transpose(1, 2)\n\n        def unshape(x):\n            \"\"\"  compute context \"\"\"\n            return x.transpose(1, 2).contiguous() \\\n                    .view(batch_size, -1, head_count * dim_per_head)\n\n        # 1) Project key, value, and query.\n        if layer_cache is not None:\n            if type == \"self\":\n                query, key, value = self.linear_query(query),\\\n                                    self.linear_keys(query),\\\n                                    self.linear_values(query)\n                key = shape(key)\n                value = shape(value)\n                device = key.device\n                if layer_cache[\"self_keys\"] is not None:\n                    key = torch.cat(\n                        (layer_cache[\"self_keys\"].to(device), key),\n                        dim=2)\n                if layer_cache[\"self_values\"] is not None:\n                    value = torch.cat(\n                        (layer_cache[\"self_values\"].to(device), value),\n                        dim=2)\n                layer_cache[\"self_keys\"] = key\n                layer_cache[\"self_values\"] = value\n            elif type == \"context\":\n                query = self.linear_query(query)\n                if layer_cache[\"memory_keys\"] is None:\n                    key, value = self.linear_keys(key),\\\n                                 self.linear_values(value)\n                    key = shape(key)\n                    value = shape(value)\n                else:\n                    key, value = layer_cache[\"memory_keys\"],\\\n                               layer_cache[\"memory_values\"]\n                layer_cache[\"memory_keys\"] = key\n                layer_cache[\"memory_values\"] = value\n        else:\n            key = self.linear_keys(key)\n            value = self.linear_values(value)\n            query = self.linear_query(query)\n            key = shape(key)\n            value = shape(value)\n\n        query = shape(query)\n\n        key_len = key.size(2)\n        query_len = query.size(2)\n\n        # 2) Calculate and scale scores.\n        query = query / math.sqrt(dim_per_head)\n        scores = torch.matmul(query, key.transpose(2, 3))\n\n        if mask is not None:\n            mask = mask.unsqueeze(1)  # [B, 1, 1, T_values]\n            scores = scores.masked_fill(mask, -1e18)\n\n        # 3) Apply attention dropout and compute context vectors.\n        attn = self.softmax(scores)\n        drop_attn = self.dropout(attn)\n        context = unshape(torch.matmul(drop_attn, value))\n\n        output = self.final_linear(context)\n        # CHECK\n        # batch_, q_len_, d_ = output.size()\n        # aeq(q_len, q_len_)\n        # aeq(batch, batch_)\n        # aeq(d, d_)\n\n        # Return one attn\n        top_attn = attn \\\n            .view(batch_size, head_count,\n                  query_len, key_len)[:, 0, :, :] \\\n            .contiguous()\n\n        return output, top_attn\n", "path": "onmt/modules/multi_headed_attn.py"}]}