self.W_Q = nn.Linear(d_model, query_key_dim * n_heads)   # Projects input to [batch x sequence x (q/k_dim*num_heads)]
self.W_K = nn.Linear(d_model, query_key_dim * n_heads)   # Projects input to [batch x sequence x (q/k_dim*num_heads)]
self.W_V = nn.Linear(d_model, value_dim * n_heads)       # Projects input to [batch x sequence x (v_dim*num_heads)]
self.proj_back = nn.Linear(value_dim * n_heads, d_model) # Projects final output of mhsa back into model dimension