Update modeling_prismatic.py

modeling_prismatic.py

@@ -437,9 +437,7 @@ class PrismaticForConditionalGeneration(PrismaticPreTrainedModel):
 
             return language_model_output
 
-        print("=" * 100)
-        print(language_model_output.attentions)
-        print("=" * 100)
+        print("Forward")
         
         
         return PrismaticCausalLMOutputWithPast(
@@ -487,6 +485,8 @@ class PrismaticForConditionalGeneration(PrismaticPreTrainedModel):
             }
         )
 
+        print("Prepare")
+
         return model_inputs
 
     # Defer to Language Model (all handle this differently, with different return types)
@@ -520,8 +520,10 @@ class OpenVLAForActionPrediction(PrismaticForConditionalGeneration):
             )
 
         # Run VLA inference
+        print("=" * 100)
         model_outputs = self.generate(input_ids, max_new_tokens=self.get_action_dim(unnorm_key), output_attentions=True, return_dict=True, **kwargs)
         print(model_outputs.shape)
+        print("=" * 100)
 
         # Extract predicted action tokens and translate into (normalized) continuous actions
         predicted_action_token_ids = generated_ids[0, -self.get_action_dim(unnorm_key) :].cpu().numpy()