Skip to content

Commit 8c86101

Browse files
authored
model: revert Gemma4 rbln_config dtype casting (#660)
1 parent 34395f5 commit 8c86101

1 file changed

Lines changed: 3 additions & 6 deletions

File tree

src/optimum/rbln/transformers/models/gemma4/modeling_gemma4.py

Lines changed: 3 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -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

Comments
 (0)