fix ckpt loading

Author: rowanz

data/groverdenoise/lm/configs/medium.json

@@ -6,7 +6,7 @@
   "hidden_act": "gelu",
   "initializer_range": 0.02,
   "intermediate_size": 4096,
-  "max_position_embeddings": 1024,
+  "max_position_embeddings": 2048,
   "num_attention_heads": 16,
   "num_hidden_layers": 24,
   "max_batch_size_per_core": 3