src/transformers/models/openai/modeling_openai.py/Attention/merge_heads
class Attention:    def merge_heads(self, x):
        x = x.permute(0, 2, 1, 3).contiguous()
        new_x_shape = x.size()[:-2] + (x.size(-2) * x.size(-1),)
        return x.view(*new_x_shape)  # in Tensorflow implementation: fct merge_states