ptrblck:

with torch.no_grad()


Hi all,
I am running into what I think is a memory leak? Not sure why my memory isn’t getting feed up while I iterate through my dataloader during inference. If anyone has ideas please let me know!
**More Details **
I noticed that the memory appears to accumulate based on the tracemalloc package at:
def __getitem__(self, index)
...
        image = (image / 255).astype(np.float32)
        mask = mask.astype(np.float32)
...

Iteration through inference DataLoader
...
    tracemalloc.start()

    snapshot1 = tracemalloc.take_snapshot()
    for (batch_idx, batch) in tqdm(enumerate(loader), total=len(loader)):
        try:
            if batch_idx % 10 == 0:
                snapshot1 = tracemalloc.take_snapshot()

            log.info(f"Running inference on {batch['image_path']}")

            batch_pred = inference_step(
                model, checkpoint_fp, batch, thresh, model_loaded=True
            )

            res = inference_ds.save_prediction_batch(
                batch_pred,
                out_folder,
                config=cfg,
                thresh=thresh,
                channel=channel,
            )
            successes.append(res)
            log.info("Inference complete")

            if batch_idx % 10 == 9:
                snapshot2 = tracemalloc.take_snapshot()
                top_stats = snapshot2.compare_to(snapshot1, "lineno")

                log.info("\n\n [ Top 10 differences ]")
                for stat in top_stats[:10]:
                    log.info(f"Snapshot difference: \n {stat}")
...