fix(framework): fix how pytorch DataContainer convert GPU tensor (#2739)

bentoml · Jul 14, 2022 · 28f0bc3 · 28f0bc3
1 parent 548312c
commit 28f0bc3
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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(),