[BUG] Training gets slower over iterations with IterativeProblem

[sangkeun00]

Training with IterativeProblem gets slower over iterations. Below is the log from the logistic regression example. The same bug was also observed with the maml example.

[2022-07-01 11:01:53] [INFO] [Problem "outer"] [Global Step 1000] [Local Step 10] loss: 0.5296119451522827
[2022-07-01 11:01:56] [INFO] [Problem "outer"] [Global Step 2000] [Local Step 20] loss: 0.3373554050922394
[2022-07-01 11:01:59] [INFO] [Problem "outer"] [Global Step 3000] [Local Step 30] loss: 0.31969979405403137
[2022-07-01 11:02:02] [INFO] [Problem "outer"] [Global Step 4000] [Local Step 40] loss: 0.31455692648887634
[2022-07-01 11:02:05] [INFO] [Problem "outer"] [Global Step 5000] [Local Step 50] loss: 0.31011053919792175
[2022-07-01 11:02:08] [INFO] [Problem "outer"] [Global Step 6000] [Local Step 60] loss: 0.3047352433204651
[2022-07-01 11:02:12] [INFO] [Problem "outer"] [Global Step 7000] [Local Step 70] loss: 0.301718533039093
[2022-07-01 11:02:15] [INFO] [Problem "outer"] [Global Step 8000] [Local Step 80] loss: 0.30068764090538025
[2022-07-01 11:02:19] [INFO] [Problem "outer"] [Global Step 9000] [Local Step 90] loss: 0.29966291785240173
[2022-07-01 11:02:22] [INFO] [Problem "outer"] [Global Step 10000] [Local Step 100] loss: 0.2992149293422699
[2022-07-01 11:02:27] [INFO] [Problem "outer"] [Global Step 11000] [Local Step 110] loss: 0.2989771068096161
[2022-07-01 11:02:31] [INFO] [Problem "outer"] [Global Step 12000] [Local Step 120] loss: 0.2986523509025574
[2022-07-01 11:02:36] [INFO] [Problem "outer"] [Global Step 13000] [Local Step 130] loss: 0.29848340153694153
[2022-07-01 11:02:42] [INFO] [Problem "outer"] [Global Step 14000] [Local Step 140] loss: 0.29845142364501953
[2022-07-01 11:02:47] [INFO] [Problem "outer"] [Global Step 15000] [Local Step 150] loss: 0.2984345257282257
[2022-07-01 11:02:53] [INFO] [Problem "outer"] [Global Step 16000] [Local Step 160] loss: 0.2983992397785187
[2022-07-01 11:03:00] [INFO] [Problem "outer"] [Global Step 17000] [Local Step 170] loss: 0.2983682453632355
[2022-07-01 11:03:08] [INFO] [Problem "outer"] [Global Step 18000] [Local Step 180] loss: 0.29832738637924194
[2022-07-01 11:03:15] [INFO] [Problem "outer"] [Global Step 19000] [Local Step 190] loss: 0.29827484488487244
[2022-07-01 11:03:23] [INFO] [Problem "outer"] [Global Step 20000] [Local Step 200] loss: 0.29820242524147034
[2022-07-01 11:03:31] [INFO] [Problem "outer"] [Global Step 21000] [Local Step 210] loss: 0.29809457063674927
[2022-07-01 11:03:38] [INFO] [Problem "outer"] [Global Step 22000] [Local Step 220] loss: 0.2979184091091156