Update modeling_prismatic.py

modeling_prismatic.py

@@ -397,6 +397,8 @@ class PrismaticForConditionalGeneration(PrismaticPreTrainedModel):
                 )
                 multimodal_labels = torch.cat([labels[:, :1], projected_patch_labels, labels[:, 1:]], dim=1)
 
+            print("Must run this")
+            
             # Dispatch to Language Model
             language_model_output = self.language_model(
                 input_ids=None,
@@ -408,7 +410,7 @@ class PrismaticForConditionalGeneration(PrismaticPreTrainedModel):
                 use_cache=use_cache,
                 output_attentions=True,
                 output_hidden_states=output_hidden_states,
-                return_dict=return_dict,
+                return_dict=True,
             )
 
         # === Otherwise =>> Assume Invalid! ===