Gen heatmap using CPU

Taxabind/TaxaBind/SatBind/clip_seg_tta.py

@@ -438,8 +438,13 @@ class ClipSegTTA:
             self.model_local.imo_encoder.to(self.device)
             self.model_local.bio_model.to(self.device)
 
+            # Move tensors to CPU before generating heatmap to avoid dtype/device mismatches
+            img_cpu = img.to(self.device) if isinstance(img, torch.Tensor) else img
+            imo_cpu = imo.to(self.device)
+            sound_cpu = sound.to(self.device) if (sound is not None and isinstance(sound, torch.Tensor)) else sound
+
             # Save final heatmap after TTA steps 
-            self.generate_heatmap(img, imo, txt, sound=sound, modality=modality)
+            self.generate_heatmap(img_cpu, imo_cpu, txt, sound=sound_cpu, modality=modality)
 
 
     def generate_query_embeds(self, img, imo, txt, sound=None, modality="image"):