fix output issue

Signed-off-by: binliu <[email protected]>

hugging_face_pipeline.py

@@ -32,7 +32,7 @@ class HuggingFacePipelineHelper:
         config_dict = kwargs.pop("config_dict", None)
         self._update_config(config, config_dict)
         model = VISTA3DModel(config)
-        model.from_pretrained(
+        model = model.from_pretrained(
             pretrained_model_name_or_path=pretrained_model_name_or_path
         )
         return VISTA3DPipeline(model, **kwargs)

vista3d_pipeline.py

@@ -433,6 +433,8 @@ class VISTA3DPipeline(Pipeline):
         return outputs
 
     def postprocess(self, outputs, **kwargs):
+        outputs[Keys.IMAGE] = outputs[Keys.IMAGE].to(self.device)
+        outputs[Keys.PRED] = outputs[Keys.PRED].to(self.device)
         for key, value in kwargs.items():
             if key not in self.POSTPROCESSING_EXTRA_ARGS:
                 logging.warning(f"Cannot set parameter {key} for postprocessing.")