Hybridize

scripts/text_generation/model/gpt.py

@@ -242,7 +242,14 @@ def __init__(self, units, vocab_size, max_length, num_layers, num_heads, dropout
                 self._final_ln = nn.LayerNorm(prefix='final_ln{}_'.format(i))
 
     def hybrid_forward(self, F, data, states=None): # pylint: disable=arguments-differ
-        """
+        """Compute
+
+        Notes
+        -----
+        If you hybridized the GPT2Model by calling net.hybridize(), you cannot
+        switch between states=None, and states=list_of_NDArray between calls to
+        the net. The hybridized model will only support the type of states used
+        during the first call after hybridization.
 
         Parameters
         ----------

scripts/text_generation/sequence_sampling.py

@@ -170,6 +170,9 @@ def generate():
                                               scorer=scorer,
                                               max_length=args.max_length - len(bos_tokens))
     inputs, begin_states = get_initial_input_state(decoder, bos_ids)
+
+    sampler._decoder.net.hybridize()  # Hybridize after we obtained the initial states
+
     # samples have shape (1, beam_size, length), scores have shape (1, beam_size)
     samples, scores, valid_lengths = sampler(inputs, begin_states)
     samples = samples[0].asnumpy()