Tensorflow Estimator: Execute an operation at a sp

Question

I built a model in Tensorflow and I'm trying to convert it into a TensorFlow Estimator. Here is an example of what I have:

train_op = tf.train.AdamOptimizer(learning_rate=lr).minimize(cost) 
saver = tf.train.Saver()
init = tf.global_variables_initializer()

assign_Wvh = pretrained_rsm.temporal_assignment(params['W'])

with tf.Session() as sess:
    sess.run(init)

    for epoch in range(epochs):
        start = time.time()
        _ = sess.run(train_op, feed_dict={x: input})
        print("%i. elapsed time: %0.2f" % (epoch, time.time() - start))

    # before saving the weights do an operation to change the weights
    # only need to perform it once at the end to avoid unecessary operations
    # that are time consuming at each iteration
    _ = sess.run(assign_Wvh)

    # save the weights
    save_path = saver.save(sess, os.path.join(weights_path, 'init.ckpt'))

I was thinking of adding this line to my model_fn (estimator):

tf.train.get_global_step() == 1000: # 1000 is my specific epoch
    do operation

But obviously I can't do that with an estimator.

Does someone know how to achieve such a thing? Knowing that I still need to save my weights that will be transformed by this last operation.