(A, V)
| 160 | |
| 161 | |
| 162 | def apply_alpha(A, V): |
| 163 | # TODO: Check this function again. |
| 164 | b, h, w, c = A.size() |
| 165 | A = A.reshape(b, h * w, c).transpose(1, 2) |
| 166 | |
| 167 | _, _, _, d = V.size() |
| 168 | V = V.reshape(b, h * w, d) |
| 169 | |
| 170 | return torch.matmul(A, V) |
| 171 | |
| 172 | |
| 173 | class SpatialAttGRU(nn.Module): |