.

Author: ZiyiTsang

verl/workers/fsdp_workers.py

@@ -479,8 +479,12 @@ def _process_multi_modal_inputs(self, data: DataProto):
                     else:
                         multi_modal_inputs = {}
 
-                    # Some image processor return with batch dim (such as glm4.1), we need to squeeze the pix_value.
-                    multi_modal_inputs['pixel_values'] = multi_modal_inputs['pixel_values'].squeeze(0)
+                    if multi_modal_inputs['pixel_values'].ndim() == 3:
+                        # Some image processor return with batch dim (such as glm4.1), we need to squeeze the pix_value.
+                        # i.e. (1,patch,pix_per_patch) -> (patch,pix_per_patch)
+                        multi_modal_inputs['pixel_values'] = multi_modal_inputs['pixel_values'].squeeze(0)
+                    
+
 
                     multi_modal_inputs_cache[index] = multi_modal_inputs