RuntimeError: The expanded size of the tensor (768) must match the existing size (512) at non-singleton dimension 1. Target sizes: [1, 768]. Tensor sizes: [1, 512]

Moe: <p>I am working on tracing a Bert based model for sentiment analysis so I can compile it using Apache TVM.</p>
<p>Here is the model:<br>
5  class BERTGRUSentiment(nn.Module):<br>
6     def <strong>init</strong>(self,<br>
7                  bert,<br>
8                  hidden_dim,<br>
9                  output_dim,<br>
10                  n_layers,<br>
11                  bidirectional,<br>
12                  dropout):<br>
13<br>
14         super().<strong>init</strong>()<br>
15<br>
16         self.bert = bert<br>
17<br>
18         embedding_dim = bert.config.to_dict()[‘hidden_size’]<br>
19<br>
20         self.rnn = nn.GRU(embedding_dim,<br>
21                           hidden_dim,<br>
22                           num_layers = n_layers,<br>
23                           bidirectional = bidirectional,<br>
24                           batch_first = True,<br>
25                           dropout = 0 if n_layers &lt; 2 else dropout)<br>
26<br>
27         self.out = nn.Linear(hidden_dim * 2 if bidirectional else hidden_dim, output_dim)<br>
28<br>
29         self.dropout = nn.Dropout(dropout)<br>
30<br>
31     def forward(self, text):<br>
32<br>
33         # text = [batch size, sent len]<br>
34<br>
35         with torch.no_grad():<br>
36             embedded = self.bert(text)[0]<br>
37<br>
38         <span class="hashtag-raw">#embedded</span> = [batch size, sent len, emb dim]<br>
39<br>
40         _, hidden = self.rnn(embedded)<br>
41<br>
42         <span class="hashtag-raw">#hidden</span> = [n layers * n directions, batch size, emb dim]<br>
43<br>
44         if self.rnn.bidirectional:<br>
45             hidden = self.dropout(torch.cat((hidden[-2,:,:], hidden[-1,:,:]), dim = 1))<br>
46         else:<br>
47             hidden = self.dropout(hidden[-1,:,:])<br>
48<br>
49         <span class="hashtag-raw">#hidden</span> = [batch size, hid dim]<br>
50<br>
51         output = self.out(hidden)<br>
52<br>
53         <span class="hashtag-raw">#output</span> = [batch size, out dim]<br>
54<br>
55         return output</p>
<p>Here is how I try to compile it:</p>
<p>5 class BERTGRUSentiment(nn.Module):<br>
6     def <strong>init</strong>(self,<br>
7                  bert,<br>
8                  hidden_dim,<br>
9                  output_dim,<br>
10                  n_layers,<br>
11                  bidirectional,<br>
12                  dropout):<br>
13<br>
14         super().<strong>init</strong>()<br>
15<br>
16         self.bert = bert<br>
17<br>
18         embedding_dim = bert.config.to_dict()[‘hidden_size’]<br>
19<br>
20         self.rnn = nn.GRU(embedding_dim,<br>
21                           hidden_dim,<br>
22                           num_layers = n_layers,<br>
23                           bidirectional = bidirectional,<br>
24                           batch_first = True,<br>
25                           dropout = 0 if n_layers &lt; 2 else dropout)<br>
26<br>
27         self.out = nn.Linear(hidden_dim * 2 if bidirectional else hidden_dim, output_dim)<br>
28<br>
29         self.dropout = nn.Dropout(dropout)<br>
30<br>
31     def forward(self, text):<br>
32<br>
33         # text = [batch size, sent len]<br>
34<br>
35         with torch.no_grad():<br>
36             embedded = self.bert(text)[0]<br>
37<br>
38         <span class="hashtag-raw">#embedded</span> = [batch size, sent len, emb dim]<br>
39<br>
40         _, hidden = self.rnn(embedded)<br>
41<br>
42         <span class="hashtag-raw">#hidden</span> = [n layers * n directions, batch size, emb dim]<br>
43<br>
44         if self.rnn.bidirectional:<br>
45             hidden = self.dropout(torch.cat((hidden[-2,:,:], hidden[-1,:,:]), dim = 1))<br>
46         else:<br>
47             hidden = self.dropout(hidden[-1,:,:])<br>
48<br>
49         <span class="hashtag-raw">#hidden</span> = [batch size, hid dim]<br>
50<br>
51         output = self.out(hidden)<br>
52<br>
53         <span class="hashtag-raw">#output</span> = [batch size, out dim]<br>
54<br>
55         return output</p>
<p>However, when I run compile.py, I get the following error;<br>
RuntimeError: The expanded size of the tensor (768) must match the existing size (512) at non-singleton dimension 1.  Target sizes: [1, 768].  Tensor sizes: [1, 512]</p>
<p>Nevertheless, I am able to successfully train the model and run in on examples. What is going wrong?</p>