diff --git a/bentoml/_internal/frameworks/common/pytorch.py b/bentoml/_internal/frameworks/common/pytorch.py index 38e5627948a..7413f3101b2 100644 --- a/bentoml/_internal/frameworks/common/pytorch.py +++ b/bentoml/_internal/frameworks/common/pytorch.py @@ -144,7 +144,7 @@ def to_payload( # pylint: disable=arguments-differ batch_dim: int = 0, plasma_db: "ext.PlasmaClient" | None = Provide[BentoMLContainer.plasma_db], ) -> Payload: - batch = batch.numpy() + batch = batch.cpu().numpy() if plasma_db: return cls.create_payload( plasma_db.put(batch).binary(),