Fix model to be properly exported to ONNX (pytorch#1144)

Co-authored-by: Brian Johnson <[email protected]>
mthrok · Sep 15, 2020 · ba6070e · ba6070e
1 parent fe33b54
commit ba6070e
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/beginner_source/transformer_tutorial.py b/beginner_source/transformer_tutorial.py
@@ -77,9 +77,9 @@ def init_weights(self):
         self.decoder.weight.data.uniform_(-initrange, initrange)
 
     def forward(self, src):
-        if self.src_mask is None or self.src_mask.size(0) != len(src):
+        if self.src_mask is None or self.src_mask.size(0) != src.size(0):
             device = src.device
-            mask = self._generate_square_subsequent_mask(len(src)).to(device)
+            mask = self._generate_square_subsequent_mask(src.size(0)).to(device)
             self.src_mask = mask
 
         src = self.encoder(src) * math.sqrt(self.ninp)