Repository: pytorch
Issue: Torch .pth model export to ONNX returned empty graph
Labels: module: onnx,triaged

## ❓ Questions and Help

### Please note that this issue tracker is not a help form and this issue will be closed.

We have a set of [listed resources available on the website](https://pytorch.org/resources). Our primary means of support is our discussion forum:

- [Discussion Forum](https://discuss.pytorch.org/)

Hi, I have a Transformer model trained and runs perfectly on my server(Ubuntu 16, tested with torch-cpu), but encountered a weird issue when trying to export the model(.pth) to ONNX, torch.export returns an empty ONNX graph without throws any error. Can anyone give an idea regarding this issue?
Following is how I convert to ONNX:
```
model = Transformer(args.token_size, args.units, args.num_heads, \
			args.num_hidden, args.num_layers, args.tokens, args.dropout, test = test).to(device)
model.eval()
checkpoint = torch.load(model_dir, map_location = device)
model_dict = checkpoint['model_state_dict']
for key in list(model_dict.keys()):
	if 'module.' in key:
		model_dict[key.replace('module.', '')] = model_dict[key]
		del model_dict[key]
model.load_state_dict(model_dict)
en_input =torch.randint(low = 0, high = args.token_size, size = args.shape,dtype = torch.long)
target_seq = torch.tensor(np.zeros((len(en_input),max_len)),dtype = torch.long)
torch.onnx.export(model,               
                  (en_input,target_seq),      
                  onnx_dir+"model.onnx",  
                  export_params=True,
                  verbose = True)
```

The output of above codes:
![image](https://user-images.githubusercontent.com/44067681/124232191-eec34b00-db43-11eb-9a50-84b2b9cce9eb.png)



**Following is my model structure (mostly the same as the native torch transformer):**


```
class Transformer(nn.Module):

    def __init__(self, ntoken, ninp, nhead, nhid, nlayers, tokens,dropout=0.2, test = False): #ninp = latent_dim, nhid = ff layers dim
        super(Transformer, self).__init__()
        self.model_type = 'Transformer'
        self.ninp = ninp
        self.ntoken = ntoken
        self.pos_encoder = PositionalEncoding(ninp, dropout)
        self.embed = nn.Embedding(ntoken, ninp)
        self.tokens = tokens
        self.test = test
        self.reversed_tokens = dict((i, char) for char, i in self.tokens.items())
        #print(self.reversed_tokens)
        encoder_layers = TransformerEncoderLayer(ninp, nhead, nhid, dropout)
        self.transformer_encoder = TransformerEncoder(encoder_layers, nlayers)
        
        
        decoder_layers = TransformerDecoderLayer(ninp, nhead, nhid, dropout)
        self.transformer_decoder = TransformerDecoder(decoder_layers, nlayers,ninp, ntoken)

        

        #self.init_weights()

    def look_ahead_mask(self, sz):
        mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1).float()
        mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
        return mask.to(device)

    def init_weights(self):
        initrange = 0.1
        self.encoder.weight.data.uniform_(-initrange, initrange)
        self.decoder.bias.data.zero_()
        self.decoder.weight.data.uniform_(-initrange, initrange)

    def forward(self, encoder_input, decoder_input):
        src = self.embed(encoder_input) # batch_size, sent_len, embed_size 
        src = torch.transpose(src,0,1) # sent_len, batch_size, embed_size
        src = self.pos_encoder(src)
        output = self.transformer_encoder(src)

        
        #print(output[50])
        if not self.test:
            tgt = self.embed(decoder_input)  # batch_size, sent_len, embed_size
            output = self.transformer_decoder(tgt, output)

            return output.to(device)

        else:
            #print(tgt.size())
            stop_condition = True
            finished = []
            j = 0
            decoded_sentence = ['' for _ in range(len(encoder_input))]
            while stop_condition:
                #print(decoder_input.shape)
                tgt = self.embed(decoder_input)  # batch_size, sent_len, embed_size
                decoder_output = self.transformer_decoder(tgt, output)

                sampled_token_index = np.argmax(decoder_output[:,j],axis = 1)
                #print(decoded_sentence, sampled_token_index)
                #print(sampled_token_index)
                #print(len(sampled_token_index))
                for i in range(len(sampled_token_index)):
                    if i in finished: continue
                    #print(sampled_token_index[i].item(), self.reversed_tokens[sampled_token_index[i]])
                    try:
                        sampled_char = self.reversed_tokens[sampled_token_index[i].item()]
                        #print(sampled_char)
                    except:
                        sampled_char = '<unk>'
                    decoded_sentence[i] += sampled_char
                    #print(decoded_sentence[i])
                    if (sampled_char == '\n' or len(decoded_sentence[i]) >= len(encoder_input[0])):
                        finished.append(i)
                        continue
                    else:
                        decoder_input[:,j+1] = sampled_token_index
                #print(finished)
                if len(finished) == len(decoder_input): break
                if j == decoder_output.size(1)-1: break
                j+=1

            return decoded_sentence
class PositionalEncoding(nn.Module):

    def __init__(self, d_model, dropout=0.1, max_len=5000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)

        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0).transpose(0, 1)
        self.register_buffer('pe', pe)

    def forward(self, x):
        x = x + self.pe[:x.size(0), :]
        return self.dropout(x)
class TransformerEncoder(Module):

    __constants__ = ['norm']

    def __init__(self, encoder_layer, num_layers, norm=None):
        super(TransformerEncoder, self).__init__()
        self.layers = _get_clones(encoder_layer, num_layers)
        self.num_layers = num_layers
        self.norm = norm

    def forward(self, src: Tensor, mask: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None) -> Tensor:

        output = src

        for mod in self.layers:
            output = mod(output, src_mask=mask, src_key_padding_mask=src_key_padding_mask)

        if self.norm is not None:
            output = self.norm(output)

        return output



class TransformerDecoder(Module):

    __constants__ = ['norm']

    def __init__(self, decoder_layer, num_layers,ninp,ntoken, norm=None):
        super(TransformerDecoder, self).__init__()
        self.layers = _get_clones(decoder_layer, num_layers)
        self.num_layers = num_layers
        self.norm = norm
        self.output_linear = nn.Linear(ninp,ntoken)

    def look_ahead_mask(self, sz):
        mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1).float()
        mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
        return mask.to(device)

    def forward(self, tgt: Tensor, memory: Tensor, tgt_mask: Optional[Tensor] = None,
                memory_mask: Optional[Tensor] = None, tgt_key_padding_mask: Optional[Tensor] = None,
                memory_key_padding_mask: Optional[Tensor] = None) -> Tensor:
        
        tgt = torch.transpose(tgt,0,1)
        tgt_mask = self.look_ahead_mask(tgt.size(0))
        output = tgt

        for mod in self.layers:
            output = mod(output, memory, tgt_mask=tgt_mask,
                         memory_mask=memory_mask,
                         tgt_key_padding_mask=tgt_key_padding_mask,
                         memory_key_padding_mask=memory_key_padding_mask)

        if self.norm is not None:
            output = self.norm(output)


        output = self.output_linear(output)
        #output = nn.Softmax(dim = -1)(output)
        output = torch.transpose(output,0,1)
        output = nn.LogSoftmax(dim = -1)(output)

        return output


class TransformerEncoderLayer(Module):

    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, activation="relu"):
        super(TransformerEncoderLayer, self).__init__()
        self.self_attn = MultiheadAttention(d_model, nhead, dropout=dropout)
        # Implementation of Feedforward model
        self.linear1 = Linear(d_model, dim_feedforward)
        self.dropout = Dropout(dropout)
        self.linear2 = Linear(dim_feedforward, d_model)

        self.norm1 = LayerNorm(d_model)
        self.norm2 = LayerNorm(d_model)
        self.dropout1 = Dropout(dropout)
        self.dropout2 = Dropout(dropout)

        self.activation = _get_activation_fn(activation)

    def __setstate__(self, state):
        if 'activation' not in state:
            state['activation'] = F.relu
        super(TransformerEncoderLayer, self).__setstate__(state)

    def forward(self, src: Tensor, src_mask: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None) -> Tensor:

        src2 = self.self_attn(src, src, src, attn_mask=src_mask,
                              key_padding_mask=src_key_padding_mask)[0]
        src = src + self.dropout1(src2)
        src = self.norm1(src)
        src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))
        src = src + self.dropout2(src2)
        src = self.norm2(src)
        return src



class TransformerDecoderLayer(Module):

    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, activation="relu"):
        super(TransformerDecoderLayer, self).__init__()
        self.self_attn = MultiheadAttention(d_model, nhead, dropout=dropout)
        self.multihead_attn = MultiheadAttention(d_model, nhead, dropout=dropout)
        # Implementation of Feedforward model
        self.linear1 = Linear(d_model, dim_feedforward)
        self.dropout = Dropout(dropout)
        self.linear2 = Linear(dim_feedforward, d_model)

        self.norm1 = LayerNorm(d_model)
        self.norm2 = LayerNorm(d_model)
        self.norm3 = LayerNorm(d_model)
        self.dropout1 = Dropout(dropout)
        self.dropout2 = Dropout(dropout)
        self.dropout3 = Dropout(dropout)

        self.activation = _get_activation_fn(activation)

    def __setstate__(self, state):
        if 'activation' not in state:
            state['activation'] = F.relu
        super(TransformerDecoderLayer, self).__setstate__(state)

    def forward(self, tgt: Tensor, memory: Tensor, tgt_mask: Optional[Tensor] = None, memory_mask: Optional[Tensor] = None,
                tgt_key_padding_mask: Optional[Tensor] = None, memory_key_padding_mask: Optional[Tensor] = None) -> Tensor:

        tgt2 = self.self_attn(tgt, tgt, tgt, attn_mask=tgt_mask,
                              key_padding_mask=tgt_key_padding_mask)[0]
        tgt = tgt + self.dropout1(tgt2)
        tgt = self.norm1(tgt)
        tgt2 = self.multihead_attn(tgt, memory, memory, attn_mask=memory_mask,
                                   key_padding_mask=memory_key_padding_mask)[0]
        tgt = tgt + self.dropout2(tgt2)
        tgt = self.norm2(tgt)
        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
        tgt = tgt + self.dropout3(tgt2)
        tgt = self.norm3(tgt)

        return tgt
def _get_clones(module, N):
    return ModuleList([copy.deepcopy(module) for i in range(N)])
def _get_activation_fn(activation):
    if activation == "relu":
        return F.relu
    elif activation == "gelu":
        return F.gelu

    raise RuntimeError("activation should be relu/gelu, not {}".format(activation))
````





cc @BowenBao @neginraoof