fix sampling error when recompute_grad = True

EleutherAI · Sep 17, 2020 · 8415e56 · 8415e56
1 parent ab942b1
commit 8415e56
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/models/gpt2/gpt2.py b/models/gpt2/gpt2.py
@@ -400,7 +400,8 @@ def fn(x):
 def model(mtf_features, other_features, params, mesh, variable_dtype, context=None):
     """A GPT style model implemented in mesh tensorflow."""
     results = {}
-    recompute_grad = params["recompute_grad"] == True  # if true, enable gradient checkpointing
+    # if true and in train mode, enable gradient checkpointing
+    recompute_grad = params["recompute_grad"] and (params["mode"] == "train") == True
     use_axial_pos_emb = params["axial_pos_emb"] != None
     no_weight_tie_emb = params["no_weight_tie"] == True
     share_parameters = exists(params["share_parameters"]) and params["share_parameters"] == True