fix: release GPU tensors after prediction

[mitchellxh]

Summary

• avoid retaining per-image probability tensors on the GPU during prediction
• detach probabilities to CPU and drop intermediate GPU tensors after each batch
• prevents large runs (~7k images) from exceeding VRAM (32GB)

Fixes #145.

[hlapp]

@mitchellxh thanks so much for your contribution!! Would you mind creating an issue that documents the problem you're seeing without this change? We can then link this PR to it.

Depending on the nature of the problem, we might also want to add a test case.

[hlapp]

@mitchellxh many thanks for posting the issue! Is the del step required, or would the following also work:

probs = self.create_probabilities(img_features, txt_features)
probs = probs.detach().cpu()

[hlapp]

And for img_features did you see that the variable going out of scope does not release the GPU memory, necessitating the explicit deletion of the object?

[mitchellxh]

@hlapp the deletion step is not needed!

img_features can be tamed by setting an appropriate batch size for the VRAM

[hlapp]

I've simplified the fix to what I understand is strictly necessary. @mitchellxh if you have time, could you check that the result should still be equivalent in outcome to your initial patch?

[Copilot]

Pull Request Overview

This PR fixes a GPU memory management issue by releasing GPU tensors after prediction to prevent VRAM accumulation during large batch processing.

• Adds explicit tensor detachment from GPU to CPU after probability computation
• Prevents GPU memory exhaustion when processing large image sets (~7k images)
• Includes explanatory comment for the memory management fix

[hlapp]

Again, thanks @mitchellxh for your contribution. I'll wait a bit for your 👍🏻 to my simplification before merging to ensure I haven't missed or misunderstood anything.

[mitchellxh]

confirmed! VRAM remains stable with your simplified fix.