Update

modeling_prismatic.py

@@ -541,7 +541,7 @@ class EmmaxForActionPrediction(PrismaticForConditionalGeneration):
         return actions, generated_ids
 
     @torch.inference_mode()
-    def generate_actions(self, image: Image, prompt_text: str, type: str, **kwargs: str) -> str:
+    def generate_actions(self, inputs, tokenizer, **kwargs: str) -> str:
         # For now, only support generation with a batch size of 1 for simplicity
         # image_transform, tokenizer = self.vision_backbone.image_transform, self.llm_backbone.tokenizer
 
@@ -557,18 +557,15 @@ class EmmaxForActionPrediction(PrismaticForConditionalGeneration):
         #     raise ValueError(f"Unsupported `pixel_values` type = {type(pixel_values)}")
 
         # Invoke super().generate --> taps into `GenerationMixin` which (redirects) to `forward()`
-        autocast_dtype = self.llm_backbone.half_precision_dtype
         # with torch.autocast("cuda", dtype=autocast_dtype, enabled=self.enable_mixed_precision_training):
-        with torch.autocast("cuda", dtype=torch.float16):
+        with torch.autocast("cuda", dtype=torch.bfloat16):
             # fmt: off
             generated_ids = self.generate(
-                input_ids=input_ids,            # Shape: [1, seq]
-                pixel_values=pixel_values,      # Shape: [1, 3, res, res] or Dict[str, Shape[1, 3, res, res]]
-                **kwargs
+                **inputs, **kwargs
             )
             # fmt: on
 
-        generated_text = tokenizer.decode(generated_ids[0, input_ids.shape[1] :], skip_special_tokens=True).strip()
+        generated_text = tokenizer.decode(generated_ids[0, inputs['input_ids'].shape[1] :], skip_special_tokens=True).strip()
 
         s = solver
         actions, reasoning = s.extract_action_policies(generated_text)
@@ -586,7 +583,7 @@ class EmmaxForActionPrediction(PrismaticForConditionalGeneration):
             )
             _actions.append(action_norm)
 
-        return _actions, generated_text
+        return _actions[0], generated_text
 
 
     @staticmethod