fix issue #2091 (#2107)

* fix EnasMutator LSTMCell forward bug fix issue #2091 * add comment Co-authored-by: QuanluZhang <[email protected]>
microsoft · Feb 29, 2020 · 03cea2b · 03cea2b
1 parent 6164207
commit 03cea2b
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/sdk/pynni/nni/nas/pytorch/enas/mutator.py b/src/sdk/pynni/nni/nas/pytorch/enas/mutator.py
@@ -23,7 +23,9 @@ def forward(self, inputs, hidden):
             curr_c, curr_h = m(inputs, (prev_c[i], prev_h[i]))
             next_c.append(curr_c)
             next_h.append(curr_h)
-            inputs = curr_h[-1]
+            # current implementation only supports batch size equals 1,
+            # but the algorithm does not necessarily have this limitation
+            inputs = curr_h[-1].view(1, -1)
         return next_c, next_h