remove deprecated back_prop kwarg from tf.while_loop

minimaxir · Oct 12, 2021 · aebb53b · aebb53b
1 parent 0cc4a02
commit aebb53b
Showing 1 changed file with 17 additions and 15 deletions.
diff --git a/gpt_2_simple/src/sample.py b/gpt_2_simple/src/sample.py
@@ -84,21 +84,23 @@ def body(past, prev, output):
         def cond(*args):
             return True
 
-        _, _, tokens = tf.while_loop(
-            cond=cond, body=body,
-            maximum_iterations=length,
-            loop_vars=[
-                context_output['presents'],
-                context[:, -1],
-                context,
-            ],
-            shape_invariants=[
-                tf.TensorShape(model.past_shape(
-                    hparams=hparams, batch_size=batch_size)),
-                tf.TensorShape([batch_size]),
-                tf.TensorShape([batch_size, None]),
-            ],
-            back_prop=False,
+        _, _, tokens = tf.nest.map_structure(
+            tf.stop_gradient,
+            tf.while_loop(
+                cond=cond,
+                body=body,
+                maximum_iterations=length,
+                loop_vars=[
+                    context_output['presents'],
+                    context[:, -1],
+                    context,
+                ],
+                shape_invariants=[
+                    tf.TensorShape(model.past_shape(hparams=hparams, batch_size=batch_size)),
+                    tf.TensorShape([batch_size]),
+                    tf.TensorShape([batch_size, None]),
+                ],
+            )
         )
 
         return tokens