@@ -711,7 +711,7 @@ def _update_rbln_config(
711711 input_infos = []
712712 for max_soft_tokens in max_soft_tokens_list :
713713 input_info = [
714- ("image_features" , [1 , max_soft_tokens , vision_hidden_size ], rbln_config . dtype ),
714+ ("image_features" , [1 , max_soft_tokens , vision_hidden_size ], "float32" ),
715715 ]
716716 input_infos .append (input_info )
717717 compile_cfgs = RBLNCompileConfig (input_info = input_infos )
@@ -814,13 +814,10 @@ def get_image_features(
814814 torch .empty (size = vision_out_0_size , dtype = self .rbln_config .vision_tower .dtype , device = "cpu" ),
815815 torch .empty (size = vision_out_1_size , dtype = torch .bool , device = "cpu" ),
816816 ]
817- projector_dtype = self .rbln_config .dtype
818- projector_out_buffer = [torch .empty (size = projector_out_size , dtype = projector_dtype , device = "cpu" )]
817+ projector_out_buffer = [torch .empty (size = projector_out_size , dtype = torch .float32 , device = "cpu" )]
819818
820819 vision_outputs = self .vision_tower (pixel_values , pixel_position_ids , out = vision_out_buffer )
821- pooler_output = self .embed_vision (
822- vision_outputs .last_hidden_state .to (projector_dtype ), out = projector_out_buffer
823- )
820+ pooler_output = self .embed_vision (vision_outputs .last_hidden_state .float (), out = projector_out_buffer )
824821 pooler_output = pooler_output [vision_outputs .pooler_mask ]
825822 return pooler_output
826823
0 commit comments