WIP

waifairer · Oct 20, 2023 · 154039e · 154039e
1 parent b70d3a7
commit 154039e
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/examples/mask_rcnn/loss.py b/examples/mask_rcnn/loss.py
@@ -211,14 +211,16 @@ def concat_box_prediction_layers(box_cls, box_regression):
   # being concatenated as well)
   if DEBUG > 0:
     print_gpu_memory("prior to cat")
+    print("len ", len(box_cls_flattened))
+    box_cls_flattened[0].realize()
+    print_gpu_memory("realize cls")
   box_cls = Tensor.cat(*box_cls_flattened, dim=1)#.reshape(-1, C)
   if DEBUG > 0:
     print_gpu_memory("after boxcls cat")
     box_cls.realize()
     print_gpu_memory("after boxcls realize") # memory taken here
     del box_cls
     gc.collect()
-    cuda.empty_cache()
     print_gpu_memory("after boxcls delete") # memory not reclaimed here!!
     del box_regression_flattened
     print_gpu_memory("after boxcls flat delete") # memory not reclaimed here!!