Performance diminish after running more than 40 steps on TPUs.

[radna0]

Performance of 40 steps vs 100 steps.

Step 1/40
  2%|██▏                                                                                    | 1/40 [00:01<01:08,  1.77s/it]Step 2/40
  5%|████▎                                                                                  | 2/40 [00:03<01:08,  1.79s/it]Step 3/40
  8%|██████▌                                                                                | 3/40 [00:03<00:40,  1.09s/it]Step 4/40
 10%|████████▋                                                                              | 4/40 [00:03<00:25,  1.40it/s]Step 5/40
 12%|██████████▉                                                                            | 5/40 [00:04<00:17,  1.99it/s]Step 6/40
 15%|█████████████                                                                          | 6/40 [00:04<00:12,  2.68it/s]Step 7/40
 18%|███████████████▏                                                                       | 7/40 [00:04<00:09,  3.43it/s]Step 8/40
 20%|█████████████████▍                                                                     | 8/40 [00:04<00:07,  4.21it/s]Step 9/40
 22%|███████████████████▌                                                                   | 9/40 [00:04<00:06,  4.97it/s]Step 10/40
 25%|█████████████████████▌                                                                | 10/40 [00:04<00:05,  5.67it/s]Step 11/40
 28%|███████████████████████▋                                                              | 11/40 [00:04<00:04,  6.27it/s]Step 12/40
 30%|█████████████████████████▊                                                            | 12/40 [00:04<00:04,  6.78it/s]Step 13/40
 32%|███████████████████████████▉                                                          | 13/40 [00:05<00:03,  7.03it/s]Step 14/40
 35%|██████████████████████████████                                                        | 14/40 [00:05<00:03,  7.35it/s]Step 15/40
 38%|████████████████████████████████▎                                                     | 15/40 [00:05<00:03,  7.60it/s]Step 16/40
 40%|██████████████████████████████████▍                                                   | 16/40 [00:05<00:03,  7.77it/s]Step 17/40
 42%|████████████████████████████████████▌                                                 | 17/40 [00:05<00:02,  7.90it/s]Step 18/40
 45%|██████████████████████████████████████▋                                               | 18/40 [00:05<00:02,  8.02it/s]Step 19/40
 48%|████████████████████████████████████████▊                                             | 19/40 [00:05<00:02,  8.06it/s]Step 20/40
 50%|███████████████████████████████████████████                                           | 20/40 [00:05<00:02,  8.12it/s]Step 21/40
 52%|█████████████████████████████████████████████▏                                        | 21/40 [00:06<00:02,  7.93it/s]Step 22/40
 55%|███████████████████████████████████████████████▎                                      | 22/40 [00:06<00:02,  7.99it/s]Step 23/40
 57%|█████████████████████████████████████████████████▍                                    | 23/40 [00:06<00:02,  8.05it/s]Step 24/40
 60%|███████████████████████████████████████████████████▌                                  | 24/40 [00:06<00:01,  8.02it/s]Step 25/40
 62%|█████████████████████████████████████████████████████▊                                | 25/40 [00:06<00:01,  8.05it/s]Step 26/40
 65%|███████████████████████████████████████████████████████▉                              | 26/40 [00:06<00:01,  8.13it/s]Step 27/40
 68%|██████████████████████████████████████████████████████████                            | 27/40 [00:06<00:01,  8.17it/s]Step 28/40
 70%|████████████████████████████████████████████████████████████▏                         | 28/40 [00:06<00:01,  8.19it/s]Step 29/40
 72%|██████████████████████████████████████████████████████████████▎                       | 29/40 [00:07<00:01,  8.22it/s]Step 30/40
 75%|████████████████████████████████████████████████████████████████▌                     | 30/40 [00:07<00:01,  8.12it/s]Step 31/40
 78%|██████████████████████████████████████████████████████████████████▋                   | 31/40 [00:07<00:01,  8.05it/s]Step 32/40
 80%|████████████████████████████████████████████████████████████████████▊                 | 32/40 [00:07<00:00,  8.10it/s]Step 33/40
 82%|██████████████████████████████████████████████████████████████████████▉               | 33/40 [00:07<00:00,  8.12it/s]Step 34/40
 85%|█████████████████████████████████████████████████████████████████████████             | 34/40 [00:07<00:00,  8.21it/s]Step 35/40
 88%|███████████████████████████████████████████████████████████████████████████▎          | 35/40 [00:09<00:03,  1.27it/s]Step 36/40
 90%|█████████████████████████████████████████████████████████████████████████████▍        | 36/40 [00:13<00:06,  1.74s/it]Step 37/40
 92%|███████████████████████████████████████████████████████████████████████████████▌      | 37/40 [00:17<00:07,  2.41s/it]Step 38/40
 95%|█████████████████████████████████████████████████████████████████████████████████▋    | 38/40 [00:21<00:05,  2.87s/it]Step 39/40
 98%|███████████████████████████████████████████████████████████████████████████████████▊  | 39/40 [00:25<00:03,  3.20s/it]Step 40/40
100%|██████████████████████████████████████████████████████████████████████████████████████| 40/40 [00:29<00:00,  1.34it/s]
Time taken for sampling:  34.37783861160278

Step 1/100
  1%|▊                                                                                     | 1/100 [00:02<03:52,  2.34s/it]Step 2/100
  2%|█▋                                                                                    | 2/100 [00:04<03:21,  2.05s/it]Step 3/100
  3%|██▌                                                                                   | 3/100 [00:04<01:57,  1.21s/it]Step 4/100
  4%|███▍                                                                                  | 4/100 [00:04<01:15,  1.28it/s]Step 5/100
  5%|████▎                                                                                 | 5/100 [00:04<00:51,  1.84it/s]Step 6/100
  6%|█████▏                                                                                | 6/100 [00:04<00:37,  2.50it/s]Step 7/100
  7%|██████                                                                                | 7/100 [00:04<00:28,  3.23it/s]Step 8/100
  8%|██████▉                                                                               | 8/100 [00:05<00:22,  4.00it/s]Step 9/100
  9%|███████▋                                                                              | 9/100 [00:05<00:19,  4.76it/s]Step 10/100
 10%|████████▌                                                                            | 10/100 [00:05<00:16,  5.45it/s]Step 11/100
 11%|█████████▎                                                                           | 11/100 [00:05<00:14,  5.99it/s]Step 12/100
 12%|██████████▏                                                                          | 12/100 [00:05<00:13,  6.51it/s]Step 13/100
 13%|███████████                                                                          | 13/100 [00:05<00:12,  6.89it/s]Step 14/100
 14%|███████████▉                                                                         | 14/100 [00:05<00:11,  7.17it/s]Step 15/100
 15%|████████████▊                                                                        | 15/100 [00:05<00:11,  7.38it/s]Step 16/100
 16%|█████████████▌                                                                       | 16/100 [00:06<00:11,  7.39it/s]Step 17/100
 17%|██████████████▍                                                                      | 17/100 [00:06<00:10,  7.57it/s]Step 18/100
 18%|███████████████▎                                                                     | 18/100 [00:06<00:10,  7.69it/s]Step 19/100
 19%|████████████████▏                                                                    | 19/100 [00:06<00:10,  7.83it/s]Step 20/100
 20%|█████████████████                                                                    | 20/100 [00:06<00:10,  7.91it/s]Step 21/100
 21%|█████████████████▊                                                                   | 21/100 [00:06<00:09,  7.97it/s]Step 22/100
 22%|██████████████████▋                                                                  | 22/100 [00:06<00:10,  7.69it/s]Step 23/100
 23%|███████████████████▌                                                                 | 23/100 [00:06<00:10,  7.44it/s]Step 24/100
 24%|████████████████████▍                                                                | 24/100 [00:07<00:09,  7.60it/s]Step 25/100
 25%|█████████████████████▎                                                               | 25/100 [00:07<00:09,  7.76it/s]Step 26/100
 26%|██████████████████████                                                               | 26/100 [00:07<00:09,  7.63it/s]Step 27/100
 27%|██████████████████████▉                                                              | 27/100 [00:07<00:09,  7.79it/s]Step 28/100
 28%|███████████████████████▊                                                             | 28/100 [00:07<00:09,  7.81it/s]Step 29/100
 29%|████████████████████████▋                                                            | 29/100 [00:07<00:09,  7.86it/s]Step 30/100
 30%|█████████████████████████▌                                                           | 30/100 [00:07<00:08,  7.95it/s]Step 31/100
 31%|██████████████████████████▎                                                          | 31/100 [00:07<00:08,  7.93it/s]Step 32/100
 32%|███████████████████████████▏                                                         | 32/100 [00:08<00:08,  8.03it/s]Step 33/100
 33%|████████████████████████████                                                         | 33/100 [00:08<00:08,  7.99it/s]Step 34/100
 34%|████████████████████████████▉                                                        | 34/100 [00:08<00:08,  8.05it/s]Step 35/100
 35%|█████████████████████████████▋                                                       | 35/100 [00:08<00:17,  3.81it/s]Step 36/100
 36%|██████████████████████████████▌                                                      | 36/100 [00:12<01:11,  1.12s/it]Step 37/100
 37%|███████████████████████████████▍                                                     | 37/100 [00:15<01:48,  1.72s/it]Step 38/100
 38%|████████████████████████████████▎                                                    | 38/100 [00:18<02:13,  2.15s/it]Step 39/100
 39%|█████████████████████████████████▏                                                   | 39/100 [00:21<02:28,  2.44s/it]Step 40/100
 40%|██████████████████████████████████                                                   | 40/100 [00:24<02:38,  2.65s/it]Step 41/100
 41%|██████████████████████████████████▊                                                  | 41/100 [00:27<02:44,  2.79s/it]Step 42/100
 42%|███████████████████████████████████▋                                                 | 42/100 [00:30<02:47,  2.89s/it]Step 43/100
 43%|████████████████████████████████████▌                                                | 43/100 [00:33<02:48,  2.96s/it]Step 44/100
 44%|█████████████████████████████████████▍                                               | 44/100 [00:37<02:48,  3.01s/it]Step 45/100
 45%|██████████████████████████████████████▎                                              | 45/100 [00:40<02:47,  3.05s/it]Step 46/100
 46%|███████████████████████████████████████                                              | 46/100 [00:43<02:45,  3.07s/it]Step 47/100
 47%|███████████████████████████████████████▉                                             | 47/100 [00:46<02:43,  3.09s/it]Step 48/100
 48%|████████████████████████████████████████▊                                            | 48/100 [00:49<02:41,  3.10s/it]Step 49/100
 49%|█████████████████████████████████████████▋                                           | 49/100 [00:52<02:38,  3.11s/it]Step 50/100
 50%|██████████████████████████████████████████▌                                          | 50/100 [00:55<02:35,  3.11s/it]Step 51/100
 51%|███████████████████████████████████████████▎                                         | 51/100 [00:58<02:32,  3.12s/it]Step 52/100
 52%|████████████████████████████████████████████▏                                        | 52/100 [01:02<02:29,  3.12s/it]Step 53/100
 53%|█████████████████████████████████████████████                                        | 53/100 [01:05<02:26,  3.12s/it]Step 54/100
 54%|█████████████████████████████████████████████▉                                       | 54/100 [01:08<02:23,  3.13s/it]Step 55/100
 55%|██████████████████████████████████████████████▊                                      | 55/100 [01:11<02:20,  3.13s/it]Step 56/100
 56%|███████████████████████████████████████████████▌                                     | 56/100 [01:14<02:17,  3.13s/it]Step 57/100
 57%|████████████████████████████████████████████████▍                                    | 57/100 [01:17<02:14,  3.13s/it]Step 58/100
 58%|█████████████████████████████████████████████████▎                                   | 58/100 [01:20<02:11,  3.13s/it]Step 59/100
 59%|██████████████████████████████████████████████████▏                                  | 59/100 [01:23<02:08,  3.13s/it]Step 60/100
 60%|███████████████████████████████████████████████████                                  | 60/100 [01:27<02:05,  3.13s/it]Step 61/100
 61%|███████████████████████████████████████████████████▊                                 | 61/100 [01:30<02:02,  3.13s/it]Step 62/100
 62%|████████████████████████████████████████████████████▋                                | 62/100 [01:33<01:58,  3.13s/it]Step 63/100
 63%|█████████████████████████████████████████████████████▌                               | 63/100 [01:36<01:55,  3.13s/it]Step 64/100
 64%|██████████████████████████████████████████████████████▍                              | 64/100 [01:39<01:52,  3.13s/it]Step 65/100
 65%|███████████████████████████████████████████████████████▎                             | 65/100 [01:42<01:49,  3.13s/it]Step 66/100
 66%|████████████████████████████████████████████████████████                             | 66/100 [01:45<01:46,  3.13s/it]Step 67/100
 67%|████████████████████████████████████████████████████████▉                            | 67/100 [01:49<01:43,  3.13s/it]Step 68/100
 68%|█████████████████████████████████████████████████████████▊                           | 68/100 [01:52<01:40,  3.13s/it]Step 69/100
 69%|██████████████████████████████████████████████████████████▋                          | 69/100 [01:55<01:36,  3.13s/it]Step 70/100
 70%|███████████████████████████████████████████████████████████▍                         | 70/100 [01:58<01:33,  3.13s/it]Step 71/100
 71%|████████████████████████████████████████████████████████████▎                        | 71/100 [02:01<01:30,  3.13s/it]Step 72/100
 72%|█████████████████████████████████████████████████████████████▏                       | 72/100 [02:04<01:27,  3.13s/it]Step 73/100
 73%|██████████████████████████████████████████████████████████████                       | 73/100 [02:07<01:24,  3.13s/it]Step 74/100
 74%|██████████████████████████████████████████████████████████████▉                      | 74/100 [02:10<01:21,  3.13s/it]Step 75/100
 75%|███████████████████████████████████████████████████████████████▊                     | 75/100 [02:14<01:18,  3.13s/it]Step 76/100
 76%|████████████████████████████████████████████████████████████████▌                    | 76/100 [02:17<01:15,  3.13s/it]Step 77/100
 77%|█████████████████████████████████████████████████████████████████▍                   | 77/100 [02:20<01:11,  3.13s/it]Step 78/100
 78%|██████████████████████████████████████████████████████████████████▎                  | 78/100 [02:23<01:08,  3.13s/it]Step 79/100
 79%|███████████████████████████████████████████████████████████████████▏                 | 79/100 [02:26<01:05,  3.13s/it]Step 80/100
 80%|████████████████████████████████████████████████████████████████████                 | 80/100 [02:29<01:02,  3.13s/it]Step 81/100
 81%|████████████████████████████████████████████████████████████████████▊                | 81/100 [02:32<00:59,  3.13s/it]Step 82/100
 82%|█████████████████████████████████████████████████████████████████████▋               | 82/100 [02:35<00:56,  3.13s/it]Step 83/100
 83%|██████████████████████████████████████████████████████████████████████▌              | 83/100 [02:39<00:53,  3.13s/it]Step 84/100
 84%|███████████████████████████████████████████████████████████████████████▍             | 84/100 [02:42<00:50,  3.13s/it]Step 85/100
 85%|████████████████████████████████████████████████████████████████████████▎            | 85/100 [02:45<00:46,  3.13s/it]Step 86/100
 86%|█████████████████████████████████████████████████████████████████████████            | 86/100 [02:48<00:43,  3.13s/it]Step 87/100
 87%|█████████████████████████████████████████████████████████████████████████▉           | 87/100 [02:51<00:40,  3.13s/it]Step 88/100
 88%|██████████████████████████████████████████████████████████████████████████▊          | 88/100 [02:54<00:37,  3.13s/it]Step 89/100
 89%|███████████████████████████████████████████████████████████████████████████▋         | 89/100 [02:57<00:34,  3.13s/it]Step 90/100
 90%|████████████████████████████████████████████████████████████████████████████▌        | 90/100 [03:00<00:31,  3.13s/it]Step 91/100
 91%|█████████████████████████████████████████████████████████████████████████████▎       | 91/100 [03:04<00:28,  3.13s/it]Step 92/100
 92%|██████████████████████████████████████████████████████████████████████████████▏      | 92/100 [03:07<00:25,  3.13s/it]Step 93/100
 93%|███████████████████████████████████████████████████████████████████████████████      | 93/100 [03:10<00:21,  3.13s/it]Step 94/100
 94%|███████████████████████████████████████████████████████████████████████████████▉     | 94/100 [03:13<00:18,  3.13s/it]Step 95/100
 95%|████████████████████████████████████████████████████████████████████████████████▊    | 95/100 [03:16<00:15,  3.13s/it]Step 96/100
 96%|█████████████████████████████████████████████████████████████████████████████████▌   | 96/100 [03:19<00:12,  3.13s/it]Step 97/100
 97%|██████████████████████████████████████████████████████████████████████████████████▍  | 97/100 [03:22<00:09,  3.13s/it]Step 98/100
 98%|███████████████████████████████████████████████████████████████████████████████████▎ | 98/100 [03:25<00:06,  3.13s/it]Step 99/100
 99%|████████████████████████████████████████████████████████████████████████████████████▏| 99/100 [03:29<00:03,  3.13s/it]Step 100/100
100%|████████████████████████████████████████████████████████████████████████████████████| 100/100 [03:32<00:00,  2.12s/it]
Time taken for sampling:  216.0481414794922