def __init__(self, num_heads, hidden_dim, attention_dropout_prob):
    super(SketchSelfAttention, self).__init__()
    if ((hidden_dim % num_heads) != 0):
        raise ValueError(('The hidden size (%d) is not a multiple of the number of attention heads (%d)' % (hidden_dim, num_heads)))
    self.hidden_dim = hidden_dim
    self.num_heads = num_heads
    self.head_dim = int((self.hidden_dim / self.num_heads))
    self.all_head_dim = (self.head_dim * self.num_heads)
    self.scale_factor = math.sqrt(self.head_dim)
    self.query = nn.Linear(self.hidden_dim, self.all_head_dim)
    self.key = nn.Linear(self.hidden_dim, self.all_head_dim)
    self.value = nn.Linear(self.hidden_dim, self.all_head_dim)
    self.dropout = nn.Dropout(attention_dropout_prob)
    self.multihead_output = None
def __init__(self, num_heads, hidden_dim, attention_dropout_prob):
    super(SketchSelfAttention, self).__init__()
    if ((hidden_dim % num_heads) != 0):
        raise ValueError(('The hidden size (%d) is not a multiple of the number of attention heads (%d)' % (hidden_dim, num_heads)))
    self.hidden_dim = hidden_dim
    self.num_heads = num_heads
    self.head_dim = int((self.hidden_dim / self.num_heads))
    self.all_head_dim = (self.head_dim * self.num_heads)
    self.scale_factor = math.sqrt(self.head_dim)
    self.query = nn.Linear(self.hidden_dim, self.all_head_dim)
    self.key = nn.Linear(self.hidden_dim, self.all_head_dim)
    self.value = nn.Linear(self.hidden_dim, self.all_head_dim)
    self.dropout = nn.Dropout(attention_dropout_prob)
    self.multihead_output = None