Add Attention to CNNs

nlp
1

I want to add an attention layer to the CNN layers.
Is this okay in Pytorch to add an attention layer like below

        input = self.conv8(input)
        input = self.batchnorm8(input)
        input = self.relu(input)
        #Attention Along Frequency and Channel Dimension  #Input Shape is [b_size X Channels X Feature X Time]
        attention_weights = self.get_attention_weights(input)
        input = input * attention_weights
        out = LSTM(input)   #Feeding to RNN then 



def get_attention_weights(self, input):
        #print("attention weights", input.shape)
        input = self.attention(input)
        #print("after attention weights", input.shape)
        input = self.attnpool(input)
        #print("after attention [pool]", input.shape)
        return input