diff --git a/src/models/embeddings.cpp b/src/models/embeddings.cpp
index 71c056aca1..6528eec3a5 100644
--- a/src/models/embeddings.cpp
+++ b/src/models/embeddings.cpp
@@ -21,7 +21,7 @@ Embeddings::Embeddings(State& state, Embeddings::Mode mode, const std::string& n
   // So only create the transient input and reuse that ortvalue for previous
   // steps in the pipeline.
   if (mode == Embeddings::Mode::Input) {
-    embeddings_ = OrtValue::CreateTensor(model_.p_device_->GetAllocator(), shape_, type_);
+    embeddings_ = OrtValue::CreateTensor(model_.p_device_inputs_->GetAllocator(), shape_, type_);
   }
 }
 
@@ -49,7 +49,7 @@ void Embeddings::UpdateSequenceLength(size_t new_length) {
     shape_[1] = new_length;
 
     if (mode_ == Embeddings::Mode::Input) {
-      embeddings_ = OrtValue::CreateTensor(model_.p_device_->GetAllocator(), shape_, type_);
+      embeddings_ = OrtValue::CreateTensor(model_.p_device_inputs_->GetAllocator(), shape_, type_);
       state_.inputs_[index_] = embeddings_.get();
     }
   }
diff --git a/src/models/multi_modal_features.cpp b/src/models/multi_modal_features.cpp
index 9577839655..881000b0a5 100644
--- a/src/models/multi_modal_features.cpp
+++ b/src/models/multi_modal_features.cpp
@@ -63,7 +63,7 @@ void MultiModalFeatures::Update(bool is_prompt) {
   // num_feature_tokens will be 0 when no image is provided
   if (!is_prompt && shape_[shape_.size() - 2] > 0) {  // if num_image_tokens > 0
     shape_[shape_.size() - 2] = 0;
-    features_ = OrtValue::CreateTensor(model_.p_device_->GetAllocator(), shape_, type_);
+    features_ = OrtValue::CreateTensor(model_.p_device_inputs_->GetAllocator(), shape_, type_);
     state_.inputs_[index_] = features_.get();
   }
 }
@@ -82,7 +82,7 @@ void MultiModalFeatures::AllocateEmptyFeatures() {
   // Skip if already allocated (avoids redundant allocation when called from
   // both EmbeddingState::SetExtraInputs and the pipeline prompt path)
   if (features_ && state_.inputs_[index_] == features_.get()) return;
-  features_ = OrtValue::CreateTensor(model_.p_device_->GetAllocator(), shape_, type_);
+  features_ = OrtValue::CreateTensor(model_.p_device_inputs_->GetAllocator(), shape_, type_);
   state_.inputs_[index_] = features_.get();
 }