MNIST training hangs in ApplyAdam kernel

[Aetf]

This happens regardless the executor is using GPU or not.

Steps to reproduce

1. run executor with EXEC_SCHED_USE_GPU=1 or EXEC_SCHED_USE_GPU=0
2. run test pytest test_mnist_tf.py

Expected

Test passes

Actual

Executor blocks waiting for kernel to finish. In the mean time the GPU utilization is zero.
The block always happens in AdamApply operation.

Logs:
GPU: exec.output.zip
CPU: exec.output.zip

[Aetf]

The bug was introduced in f1b9331, found by git bisect

[Aetf]

Fixed in 61fc138