Update modeling_molmo_point.py

modeling_molmo_point.py

@@ -1497,7 +1497,7 @@ class MolmoPointModel(MolmoPointPreTrainedModel):
                 assert last_predicted_patch_id is not None, "Patch should always be generated before a subpatch"
                 for_patches = (last_predicted_patch_id.view(batch_size) + image_token_offset)[input_subpatch_ids.view(batch_size) >= 0]
                 vit_features_to_embed = vit_features_flat[for_patches, input_subpatch_ids]
-                x.view(-1, dim)[is_subpatch.view(-1)] = self.build_vit_embedding(vit_features_to_embed).to(device=device)
+                x.view(-1, dim)[is_subpatch.view(-1)] = self.build_vit_embedding(vit_features_to_embed).to(device=device, dtype=x.dtype)
 
         # shape: (batch_size, seq_len, d_model)
         x = self.transformer.emb_drop(x)  # type: ignore