Solved by matheo_r in post #4 
              
              
                Hi @jack_ma 
Sorry for the late answer. You could write a custom class using MultiHeadAttention module as its parent: 
import torch
import torch.nn as nn

class myAttentionModule(nn.MultiheadAttention):
	def __init__(self, embed_dim, num_heads):
		super(myAttentionModule, self).__init__(embed_dim, n…