rename x to idx for consistency with dynamic shapes

pytorch · Apr 6, 2024 · fbb6130 · fbb6130
1 parent 8dcb622
commit fbb6130
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/export.py b/export.py
@@ -46,10 +46,10 @@ def __init__(self, model, device, max_seq_length=1024):
         self.model = model
         # init model here if necessary
 
-    def forward(self, x, input_pos):
+    def forward(self, idx, input_pos):
         # input_pos: [B, 1]
         assert input_pos.shape[-1] == 1
-        logits = self.model(x, input_pos)
+        logits = self.model(idx, input_pos)
         return logits  # sample(logits, **sampling_kwargs)