fix CityscapesSOTA memory leak bug (#845)

PaddlePaddle · Feb 23, 2021 · a5902ab · a5902ab
1 parent 6d6ae57
commit a5902ab
Showing 1 changed file with 3 additions and 5 deletions.
diff --git a/contrib/CityscapesSOTA/scripts/train.py b/contrib/CityscapesSOTA/scripts/train.py
@@ -163,18 +163,16 @@ def train(model,
             model.clear_gradients()
             avg_loss += loss.numpy()[0]
             if not avg_loss_list:
-                avg_loss_list = [l for l in loss_list]
+                avg_loss_list = [l.numpy() for l in loss_list]
             else:
                 for i in range(len(loss_list)):
-                    avg_loss_list[i] += loss_list[i]
+                    avg_loss_list[i] += loss_list[i].numpy()
             batch_cost_averager.record(
                 time.time() - batch_start, num_samples=batch_size)
 
             if (iter) % log_iters == 0 and local_rank == 0:
                 avg_loss /= log_iters
-                avg_loss_list = [
-                    l.numpy()[0] / log_iters for l in avg_loss_list
-                ]
+                avg_loss_list = [l[0] / log_iters for l in avg_loss_list]
                 remain_iters = iters - iter
                 avg_train_batch_cost = batch_cost_averager.get_average()
                 avg_train_reader_cost = reader_cost_averager.get_average()