update

demo_hf.py

@@ -65,6 +65,7 @@ def demo_fn(cfg: DictConfig, model) -> None:
         if isinstance(y_hat[key], torch.Tensor):
             y_hat[key] = y_hat[key].cpu().numpy()
 
+    torch.cuda.empty_cache()
     return y_hat
 
 

vggt/models/vggt.py

@@ -103,6 +103,7 @@ class VGGT(nn.Module):
         aggregated_tokens_list, _, patch_start_idx = self.aggregator(images, batch=batch)
 
 
+        torch.cuda.empty_cache()
         # Pose branch
         # TODO check pose encoding conversion  # Removed TODO
         # loss = 0