File tree Expand file tree Collapse file tree
python/sglang/srt/multimodal/processors Expand file tree Collapse file tree Original file line number Diff line number Diff line change @@ -221,6 +221,13 @@ def process_mm_data(
221221 return_tensors = "pt" ,
222222 ** kwargs ,
223223 )
224+ # move feature tensors to cpu
225+ for feature_name in self .FEATURE_NAMES :
226+ if feature_name in result and isinstance (
227+ result [feature_name ], torch .Tensor
228+ ):
229+ result [feature_name ] = result [feature_name ].to ("cpu" )
230+
224231 return result
225232
226233 @abstractmethod
@@ -623,19 +630,4 @@ def process_and_combine_mm_data(
623630 mm_token_id = mm_token_id ,
624631 )
625632
626- # post-process
627- for item in all_collected_items :
628- # replace the feature tensor with a proxy
629- if isinstance (item .feature , torch .Tensor ) and item .feature .is_cuda :
630- item .feature = TransportProxyTensor (
631- transport_mode = self .transport_mode , data = item .feature
632- )
633- elif (
634- isinstance (item .precomputed_embeddings , torch .Tensor )
635- and item .precomputed_embeddings .is_cuda
636- ):
637- item .precomputed_embeddings = TransportProxyTensor (
638- transport_mode = self .transport_mode , data = item .precomputed_embeddings
639- )
640-
641633 return all_collected_items , input_ids , ret
You can’t perform that action at this time.
0 commit comments