Save checkpoint 0 for TensorBoard visualization
@@ -117,6 +117,9 @@ class Trainer(object):
last_step = 0
with self._summary_writer.as_default():
+ if self._optimizer.iterations.numpy() == 0:
+ self._checkpoint.save(0)
+
for i, (loss, num_words) in enumerate(_forward()):
if i == 0 or (i + 1) % accum_steps == 0:
_step()
