From 560420fe717818ab61f72dfd2e8ac1819c0640e0 Mon Sep 17 00:00:00 2001 From: Shah Jawad Islam Date: Mon, 7 Jul 2025 16:33:36 +0600 Subject: [PATCH] Fix CUDA tensor serialization error in predictions saving - Convert CUDA tensors to CPU before saving with np.savez - Handle both individual tensors and lists of tensors - Ensures compatibility with numpy serialization Fixes issue where predictions containing CUDA tensors would fail to save due to numpy's inability to serialize CUDA tensors directly. --- demo_gradio.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/demo_gradio.py b/demo_gradio.py index 466a5ff4..0c1d119e 100644 --- a/demo_gradio.py +++ b/demo_gradio.py @@ -212,6 +212,14 @@ def gradio_demo( with torch.no_grad(): predictions = run_model(target_dir, model) + + # Convert CUDA tensors to CPU before saving + for key, value in predictions.items(): + if isinstance(value, list): + predictions[key] = [item.cpu() if hasattr(item, 'is_cuda') and item.is_cuda else item for item in value] + elif hasattr(value, 'is_cuda') and value.is_cuda: + predictions[key] = value.cpu() + # Save predictions prediction_save_path = os.path.join(target_dir, "predictions.npz") np.savez(prediction_save_path, **predictions)