Attempt to deallocate before regenerating heatmap:

Taxabind/TaxaBind/SatBind/clip_seg_tta.py

@@ -430,9 +430,6 @@ class ClipSegTTA:
                         species_name=self.species_name
                     )
 
-            # Save final heatmap after TTA steps 
-            self.generate_heatmap(img, imo, txt, sound=sound, modality=modality)
-
             ## NOTE: Added due to app.py (to allocate to GPU only when needed on HF)
             # if self.device.type == "cuda":
             print("Deallocating models from GPU...")
@@ -441,6 +438,10 @@ class ClipSegTTA:
             self.model_local.imo_encoder.to(self.device)
             self.model_local.bio_model.to(self.device)
 
+            # Save final heatmap after TTA steps 
+            self.generate_heatmap(img, imo, txt, sound=sound, modality=modality)
+
+
     def generate_query_embeds(self, img, imo, txt, sound=None, modality="image"):
 
         # Query Embeds