File size: 949 Bytes
4c3d957 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 | diff --git a/trellis/pipelines/trellis_image_to_3d.py b/trellis/pipelines/trellis_image_to_3d.py
index a3a7e46..49dff10 100644
--- a/trellis/pipelines/trellis_image_to_3d.py
+++ b/trellis/pipelines/trellis_image_to_3d.py
@@ -985,8 +985,10 @@ class TrellisVGGTTo3DPipeline(TrellisImageTo3DPipeline):
else:
raise ValueError(f"Unsupported type of image: {type(image)}")
- if self.low_vram:
- self.VGGT_model.to(self.device)
+ # FIX: always move VGGT to the pipeline device before use. Previously this
+ # was gated on self.low_vram, so with low_vram=False the model stayed on CPU
+ # while images were on cuda -> "tensors on cuda:0 and cpu" device mismatch.
+ self.VGGT_model.to(self.device)
with torch.no_grad():
with torch.cuda.amp.autocast(dtype=self.VGGT_dtype):
# Predict attributes including cameras, depth maps, and point maps.
|