Repository: bentoml/BentoML
Pull request: doc: add example of adaptive batching for collection parameter types
Signed-off-by: Frost Ming <me@frostming.com>

## What does this PR address?

<!--
Thanks for sending a pull request!

Congrats for making it this far! Here's a 🍱 for you. There are still a few steps ahead.

Please make sure to read the contribution guidelines, then fill out the blanks below before requesting a code review.

Name your Pull Request with one of the following prefixes, e.g. "feat: add support for PyTorch", to indicate the type of changes proposed. This is based on the [Conventional Commits specification](https://www.conventionalcommits.org/en/v1.0.0/#summary).
  - feat: (new feature for the user, not a new feature for build script)
  - fix: (bug fix for the user, not a fix to a build script)
  - docs: (changes to the documentation)
  - style: (formatting, missing semicolons, etc; no production code change)
  - refactor: (refactoring production code, eg. renaming a variable)
  - perf: (code changes that improve performance)
  - test: (adding missing tests, refactoring tests; no production code change)
  - chore: (updating grunt tasks etc; no production code change)
  - build: (changes that affect the build system or external dependencies)
  - ci: (changes to configuration files and scripts)
  - revert: (reverts a previous commit)

Describe your changes in detail. Attach screenshots here if appropriate.

Once you're done with this, someone from BentoML team or community member will help review your PR (see "Who can help review?" section for potential reviewers.). If no one has reviewed your PR after a week have passed, don't hesitate to post a new comment and ping @-the same person. Notifications sometimes get lost 🥲.
-->

<!-- Remove if not applicable -->

Fixes #3805

## Before submitting:

<!--- Go over all the following points, and put an `x` in all the boxes that apply. -->
<!--- If you're unsure about any of these, don't hesitate to ask. We're here to help! -->
<!--- If you plan to update documentation or tests in follow-up, please note -->

- [x] Does the Pull Request follow [Conventional Commits specification](https://www.conventionalcommits.org/en/v1.0.0/#summary) naming? Here are [GitHub's
      guide](https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/proposing-changes-to-your-work-with-pull-requests/creating-a-pull-request) on how to create a pull request.
- [x] Does the code follow BentoML's code style, both `make format` and `make lint` script have passed ([instructions](https://github.com/bentoml/BentoML/blob/main/DEVELOPMENT.md#style-check-auto-formatting-type-checking))?
- [x] Did you read through [contribution guidelines](https://github.com/bentoml/BentoML/blob/main/CONTRIBUTING.md#ways-to-contribute) and follow [development guidelines](https://github.com/bentoml/BentoML/blob/main/DEVELOPMENT.md#start-developing)?
- [x] Did your changes require updates to the documentation? Have you updated
      those accordingly? Here are [documentation guidelines](https://github.com/bentoml/BentoML/tree/main/docs) and [tips on writting docs](https://github.com/bentoml/BentoML/tree/main/docs#writing-documentation).
- [ ] Did you write tests to cover your changes?


Problem statement:
feature: Support for Dict of Tensors in production mode
### Feature request

When running `serve --production`, the invocation to `runner.async_run()` goes through an instance of `DataContainer` according to the input type. If you have a PyTorch model and passes as input a `torch.Tensor`, it will be serialized using `PyTorchTensorContainer`, sent to the runner server and then be deserialized there (and possibly having batches combined and so on). It also works if you invoke it using a numpy array or a pandas dataframe. It will be serialized and then converted into a Tensor later. However, it only works if your PyTorch model receives one or more tensors (eg. forward(tensor1: torch.Tensor, tensor2: torch.Tensor)`. If your model receives a Dict of Tensors, it doesn't work. And it's very common for models to receive this kind of input (BERT for example, that receives a dict with 'input_ids', 'token_type_ids', and 'attention_mask').

When you try to invoke the model with a dict, it will work outside of production mode. But in production mode, it will try to use the `DefaultContainer`, but it fails to deserialize the dict correctly and you get an error from the runner server like:

```
Traceback (most recent call last):
  File "/home/fernando/.local/share/virtualenvs/page-tagging-svc-6EE2Yu73/lib/python3.10/site-packages/uvicorn/protocols/http/h11_impl.py", line 429, in run_asgi
    result = await app(  # type: ignore[func-returns-value]
  File "/home/fernando/.local/share/virtualenvs/page-tagging-svc-6EE2Yu73/lib/python3.10/site-packages/uvicorn/middleware/proxy_headers.py", line 78, in __call__
    return await self.app(scope, receive, send)
  File "/home/fernando/.local/share/virtualenvs/page-tagging-svc-6EE2Yu73/lib/python3.10/site-packages/uvicorn/middleware/message_logger.py", line 86, in __call__
    raise exc from None
  File "/home/fernando/.local/share/virtualenvs/page-tagging-svc-6EE2Yu73/lib/python3.10/site-packages/uvicorn/middleware/message_logger.py", line 82, in __call__
    await self.app(scope, inner_receive, inner_send)
  File "/home/fernando/.local/share/virtualenvs/page-tagging-svc-6EE2Yu73/lib/python3.10/site-packages/starlette/applications.py", line 118, in __call__
    await self.middleware_stack(scope, receive, send)
  File "/home/fernando/.local/share/virtualenvs/page-tagging-svc-6EE2Yu73/lib/python3.10/site-packages/starlette/middleware/errors.py", line 184, in __call__
    raise exc
  File "/home/fernando/.local/share/virtualenvs/page-tagging-svc-6EE2Yu73/lib/python3.10/site-packages/starlette/middleware/errors.py", line 162, in __call__
    await self.app(scope, receive, _send)
  File "/home/fernando/.local/share/virtualenvs/page-tagging-svc-6EE2Yu73/lib/python3.10/site-packages/opentelemetry/instrumentation/asgi/__init__.py", line 581, in __call__
    await self.app(scope, otel_receive, otel_send)
  File "/home/fernando/.local/share/virtualenvs/page-tagging-svc-6EE2Yu73/lib/python3.10/site-packages/bentoml/_internal/server/http/instruments.py", line 293, in __call__
    await self.app(scope, receive, wrapped_send)
  File "/home/fernando/.local/share/virtualenvs/page-tagging-svc-6EE2Yu73/lib/python3.10/site-packages/bentoml/_internal/server/http/access.py", line 126, in __call__
    await self.app(scope, receive, wrapped_send)
  File "/home/fernando/.local/share/virtualenvs/page-tagging-svc-6EE2Yu73/lib/python3.10/site-packages/starlette/middleware/exceptions.py", line 79, in __call__
    raise exc
  File "/home/fernando/.local/share/virtualenvs/page-tagging-svc-6EE2Yu73/lib/python3.10/site-packages/starlette/middleware/exceptions.py", line 68, in __call__
    await self.app(scope, receive, sender)
  File "/home/fernando/.local/share/virtualenvs/page-tagging-svc-6EE2Yu73/lib/python3.10/site-packages/starlette/routing.py", line 706, in __call__
    await route.handle(scope, receive, send)
  File "/home/fernando/.local/share/virtualenvs/page-tagging-svc-6EE2Yu73/lib/python3.10/site-packages/starlette/routing.py", line 276, in handle
    await self.app(scope, receive, send)
  File "/home/fernando/.local/share/virtualenvs/page-tagging-svc-6EE2Yu73/lib/python3.10/site-packages/starlette/routing.py", line 66, in app
    response = await func(request)
  File "/home/fernando/.local/share/virtualenvs/page-tagging-svc-6EE2Yu73/lib/python3.10/site-packages/bentoml/_internal/server/runner_app.py", line 276, in _request_handler
    payload = await infer(params)
  File "/home/fernando/.local/share/virtualenvs/page-tagging-svc-6EE2Yu73/lib/python3.10/site-packages/bentoml/_internal/marshal/dispatcher.py", line 180, in _func
    raise r
  File "/home/fernando/.local/share/virtualenvs/page-tagging-svc-6EE2Yu73/lib/python3.10/site-packages/bentoml/_internal/marshal/dispatcher.py", line 444, in outbound_call
    outputs = await self.callback(
  File "/home/fernando/.local/share/virtualenvs/page-tagging-svc-6EE2Yu73/lib/python3.10/site-packages/bentoml/_internal/server/runner_app.py", line 237, in infer_batch
    batch_ret = await runner_method.async_run(
  File "/home/fernando/.local/share/virtualenvs/page-tagging-svc-6EE2Yu73/lib/python3.10/site-packages/bentoml/_internal/runner/runner.py", line 55, in async_run
    return await self.runner._runner_handle.async_run_method(self, *args, **kwargs)
  File "/home/fernando/.local/share/virtualenvs/page-tagging-svc-6EE2Yu73/lib/python3.10/site-packages/bentoml/_internal/runner/runner_handle/local.py", line 59, in async_run_method
    return await anyio.to_thread.run_sync(
  File "/home/fernando/.local/share/virtualenvs/page-tagging-svc-6EE2Yu73/lib/python3.10/site-packages/anyio/to_thread.py", line 31, in run_sync
    return await get_asynclib().run_sync_in_worker_thread(
  File "/home/fernando/.local/share/virtualenvs/page-tagging-svc-6EE2Yu73/lib/python3.10/site-packages/anyio/_backends/_asyncio.py", line 937, in run_sync_in_worker_thread
    return await future
  File "/home/fernando/.local/share/virtualenvs/page-tagging-svc-6EE2Yu73/lib/python3.10/site-packages/anyio/_backends/_asyncio.py", line 867, in run
    result = context.run(func, *args)
  File "/home/fernando/.local/share/virtualenvs/page-tagging-svc-6EE2Yu73/lib/python3.10/site-packages/bentoml/_internal/runner/runnable.py", line 139, in method
    return self.func(obj, *args, **kwargs)
  File "/home/fernando/.local/share/virtualenvs/page-tagging-svc-6EE2Yu73/lib/python3.10/site-packages/bentoml/_internal/frameworks/common/pytorch.py", line 107, in _run
    return getattr(self.model, method_name)(
  File "/home/fernando/.local/share/virtualenvs/page-tagging-svc-6EE2Yu73/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/fernando/bentoml/bentos/page-tagging-svc/ibwdpixeoopskldn/src/src/torch/model/transformer.py", line 106, in forward
    input_ids=batch["input_ids"],
TypeError: list indices must be integers or slices, not str
```

The dict somehow becomes a list of the keys of the dict.

As suggestion, BentoML could have a `DictContainer`, which would use `AutoContainer` for each value inside the dictionary. The `AutoContainer` would find the correct `DataContainer` for each type (`PyTorchTensorContainer` for tensors, for example), that would be able to handle the serialization and deserialization, combining batches and so on.

### Motivation

I need to use BentoML with a model that receives a `Dict[str, torch.Tensor]`.

### Other

I have the following workaround in place for now. It's worth noting that it only handles Dict of Tensors though.

```python
class PyTorchTensorDictContainer(DataContainer[Dict[str, torch.Tensor], torch.Tensor]):
    @classmethod
    def batches_to_batch(
        cls,
        batches: t.Sequence[Dict[str, torch.Tensor]],
        batch_dim: int = 0,
    ) -> t.Tuple[Dict[str, torch.Tensor], list[int]]:
        batch = {}
        for key in batches[0].keys():
            batch[key] = torch.cat(
                tuple(subbatch[key] for subbatch in batches), dim=batch_dim
            )
        indices = list(
            itertools.accumulate(subbatch[key].shape[batch_dim] for subbatch in batches)
        )
        indices = [0] + indices
        return batch, indices

    @classmethod
    def batch_to_batches(
        cls,
        batch: Dict[str, torch.Tensor],
        indices: t.Sequence[int],
        batch_dim: int = 0,
    ) -> t.List[Dict[str, torch.Tensor]]:
        sizes = [indices[i] - indices[i - 1] for i in range(1, len(indices))]
        output: list[Dict[str, torch.Tensor]] = [{} for _ in range(len(sizes))]
        for key in batch.keys():
            split_tensor = torch.split(batch[key], sizes, dim=batch_dim)
            for i, tensor in enumerate(split_tensor):
                output[i][key] = tensor
        return output

    @classmethod
    @inject
    def to_payload(  # pylint: disable=arguments-differ
        cls,
        batch: Dict[str, torch.Tensor],
        batch_dim: int = 0,
    ) -> Payload:
        batch = batch.copy()
        for key in batch.keys():
            batch[key] = batch[key].cpu().numpy()

        return cls.create_payload(
            pickle.dumps(batch),
            batch_size=batch[key].shape[batch_dim],
            meta={"plasma": False},
        )

    @classmethod
    @inject
    def from_payload(  # pylint: disable=arguments-differ
        cls,
        payload: Payload,
    ) -> Dict[str, torch.Tensor]:
        ret: Dict[str, np.ndarray] = pickle.loads(payload.data)
        output = {}
        for key in ret.keys():
            output[key] = torch.from_numpy(ret[key]).requires_grad_(False)
        return output

    @classmethod
    @inject
    def batch_to_payloads(  # pylint: disable=arguments-differ
        cls,
        batch: Dict[str, torch.Tensor],
        indices: t.Sequence[int],
        batch_dim: int = 0,
    ) -> t.List[Payload]:
        batches = cls.batch_to_batches(batch, indices, batch_dim)
        payloads = [cls.to_payload(i, batch_dim=batch_dim) for i in batches]
        return payloads

    @classmethod
    @inject
    def from_batch_payloads(  # pylint: disable=arguments-differ
        cls,
        payloads: t.Sequence[Payload],
        batch_dim: int = 0,
    ) -> t.Tuple[Dict[str, torch.Tensor], list[int]]:
        batches = [cls.from_payload(payload) for payload in payloads]
        return cls.batches_to_batch(batches, batch_dim)

def make_pytorch_runnable_method(
    method_name: str,
    partial_kwargs: dict[str, t.Any] | None = None,
) -> t.Callable[..., torch.Tensor]:
    if partial_kwargs is None:
        partial_kwargs = {}

    def _run(
        self: PytorchModelRunnable,
        *args: t.Dict[str, torch.Tensor],
        **kwargs,
    ) -> torch.Tensor:
        params = Params(*args, **kwargs)

        def _mapping(item: t.Dict[str, torch.Tensor]) -> t.Dict[str, torch.Tensor]:
            item = item.copy()
            for key, tensor in item.items():
                item[key] = tensor.to(self.device_id)

            return item

        with inference_mode_ctx():
            params = params.map(_mapping)
            return getattr(self.model, method_name)(
                *params.args,
                **dict(partial_kwargs, **params.kwargs),
            )

    return _run

def get_runnable(bento_model: Model):
    partial_kwargs: t.Dict[str, t.Any] = bento_model.info.options.partial_kwargs  # type: ignore

    runnable_class: type[PytorchModelRunnable] = partial_class(
        PytorchModelRunnable,
        bento_model=bento_model,
        loader=load_model,
    )
    for method_name, options in bento_model.info.signatures.items():
        method_partial_kwargs = partial_kwargs.get(method_name)
        runnable_class.add_method(
            make_pytorch_runnable_method(method_name, method_partial_kwargs),
            name=method_name,
            batchable=options.batchable,
            batch_dim=options.batch_dim,
            input_spec=options.input_spec,
            output_spec=options.output_spec,
        )
    return runnable_class
```

Then, instead of creating the runner with `model.to_runner()`, I use:

```python
runner = Runner(
    get_runnable(model),
    name=model.tag.name,
    models=[model],
)
```

Hints from issue discussion:
Hi @fernandocamargoai , Thanks for your feedback. I think the problem you met is caused by our current batch combining method for dictionary (which simply group dictionaries to list) does not match how pytorch's batch is combined (which should combine tensors inside the dictionary). We are working to improve this area now. Could you also paste your `service.py`. That will help us to learn more about our user's workflow. Thanks!
Sure. Here is my `service.py`:

```python
import logging

import bentoml
import torch
from bentoml import Runner
from bentoml._internal.runner.container import DataContainerRegistry
from bentoml.exceptions import UnprocessableEntity
from bentoml.io import JSON

from src.deployment.torch.bentoml_ext import PyTorchTensorDictContainer, get_runnable
from src.content_fetching import fetch
from src.schema import Input, Output, Tag, ApiInput, ApiOutput

model = bentoml.pytorch.get("page_tagging_torch_model:latest")
runner = Runner(
    get_runnable(model),
    name=model.tag.name,
    models=[model],
)
preprocessing_fn = model.custom_objects["preprocessing_fn"]
label_names = model.custom_objects["label_names"]

svc = bentoml.Service(name="page-tagging-svc", runners=[runner])

logger = logging.getLogger(__name__)

DataContainerRegistry.register_container(
    dict,
    dict,
    PyTorchTensorDictContainer,
)


@svc.api(input=JSON(pydantic_model=ApiInput), output=JSON(pydantic_model=ApiOutput))
async def predict(api_input: ApiInput) -> ApiOutput:
    for i, instance in enumerate(api_input.instances):
        if instance.html is None or instance.text is None:
            logger.debug("Fetching url: %s", instance.url)
            api_input.instances[i] = fetch(instance.url)
            if api_input.instances[i] is None:
                raise UnprocessableEntity("Cannot fetch url: %s" % instance.url)

    batch = {
        "url": [instance.url for instance in api_input.instances],
        "title": [instance.title for instance in api_input.instances],
        "html": [instance.html for instance in api_input.instances],
        "text": [instance.text for instance in api_input.instances],
    }

    if preprocessing_fn is not None:
        batch = preprocessing_fn(batch)

    logits_tensor = await runner.async_run(batch)
    probas = torch.sigmoid(logits_tensor).cpu().numpy()

    return ApiOutput(
        predictions=[
            Output(
                tags=[
                    Tag(label=label_names[j], proba=probas[i, j])
                    for j in range(len(label_names))
                    if probas[i, j] > 0.5
                ]
            )
            for i in range(len(api_input.instances))
        ]
    )
```
@fernandocamargoai Thanks for your report.

I have updated with an example of batching parameters of collection types(#3873 ). It is achieved in a simpler way without involving data containers. Would you please take a look and tell me if that works for you?

Commit message: Update docs/source/guides/batching.rst

Co-authored-by: Aaron Pham <29749331+aarnphm@users.noreply.github.com>

Patch to review:
diff --git a/docs/source/guides/batching.rst b/docs/source/guides/batching.rst
index 8623c5e09a9..8d1284be3d2 100644
--- a/docs/source/guides/batching.rst
+++ b/docs/source/guides/batching.rst
@@ -88,3 +88,52 @@ Error handling
 If adaptive batching cannot keep up with rate of the incoming requests while satisfying the max
 latency configuration, HTTP 503 Service Unavailable is returned. To workaround this error, consider
 relaxing the max latency requirement and further scaling the underlying hardware resources.
+
+Custom Batching
+---------------
+
+Currently, adaptive batching is only effective for certain types of parameters including(non-exhaustive) ``numpy.ndarray``, ``pandas.Series`` and framework-specific types such as ``torch.Tensor``.
+Batch parameters of other types are simply collected into a list and passed to the inference function. If your model accepts parameters that are not batchable by default,
+you can achieve adaptive batching by wrapping the inference function with a :ref:`Runner <concepts/runner:Custom Runner>`.
+
+We will demonstrate this with a PyTorch example which accepts a dictionary of ``torch.Tensor`` as input.
+
+.. code-block:: python
+    :caption: `service.py`
+
+    import bentoml
+    import torch
+
+    class MyRunnable(bentoml.Runnable):
+        def __init__(self):
+            super().__init__()
+            # if torch.cuda.device_count():
+            if torch.cuda.is_available():
+                self.device_id = "cuda"
+                # by default, torch.FloatTensor will be used on CPU.
+                torch.set_default_tensor_type("torch.cuda.FloatTensor")
+            else:
+                self.device_id = "cpu"
+            self.model = bentoml.pytorch.load_model(device_id=self.device_id)
+            # We want to turn off dropout and batchnorm when running inference.
+            self.model.train(False)
+
+        @bentoml.Runnable.method(batchable=True)
+        def __call__(self, **batch):
+            # Our model accepts a dictionary as input, but we use ``**``
+            # to deconstruct it to keyword arguments, so they can be batched rightly.
+            # Move the tensors to the target device and pass it to the model as a dict.
+            batch = {k: tensor.to(self.device_id) for k, tensor in batch.items()}
+            with torch.inference_mode():
+                return self.model(batch)
+
+    model = bentoml.pytorch.get("my_pytorch_model")
+    # Build the runner from the runnable manually, instead of calling model.to_runner() method
+    runner = bentoml.Runner(MyRunnable)
+    svc = bentoml.Service(name="my_ml_service", runners=[runner])
+
+    # Define an API endpoint for the inference
+    @svc.api(input=bentoml.io.JSON(), output=bentoml.io.JSON())
+    async def predict(api_input):
+        # Deconstruct the input dictionary to keyword arguments for batching
+        return await runner.async_run(**api_input)


Code review comments:
diff_hunk: @@ -88,3 +88,52 @@ Error handling
 If adaptive batching cannot keep up with rate of the incoming requests while satisfying the max
 latency configuration, HTTP 503 Service Unavailable is returned. To workaround this error, consider
 relaxing the max latency requirement and further scaling the underlying hardware resources.
+
+Custom Batching
+---------------
+
+Currently, adaptive batching is only effective for certain types of parameters including(non-exhaustive) ``numpy.ndarray``, ``pandas.Series`` and framework-specific types such as ``torch.Tensor``.
+Batch parameters of other types are simply collected into a list and passed to the inference function. If your model accepts parameters that are not batchable by default,
+you can achieve adaptive batching by wrapping the inference function with a :ref:`Runner <concepts/runner:Custom Runner>`.
+
+We will demonstrate this with a PyTorch example which accepts a dictionary of ``torch.Tensor`` as input.
+
+.. code-block:: python
+    :caption: `service.py`
+
+    import bentoml
+    import torch
+
+    class MyRunnable(bentoml.Runnable):
+        def __init__(self):
+            super().__init__()
+            # if torch.cuda.device_count():
+            if torch.cuda.is_available():
+                self.device_id = "cuda"
+                # by default, torch.FloatTensor will be used on CPU.
+                torch.set_default_tensor_type("torch.cuda.FloatTensor")
+            else:
+                self.device_id = "cpu"
+            self.model = bentoml.pytorch.load_model(device_id=self.device_id)
path: docs/source/guides/batching.rst
text: @user1:
```suggestion
            self.model = bentoml.pytorch.load_model("my_pytorch_model", device_id=self.device_id)
```

diff_hunk: @@ -88,3 +88,52 @@ Error handling
 If adaptive batching cannot keep up with rate of the incoming requests while satisfying the max
 latency configuration, HTTP 503 Service Unavailable is returned. To workaround this error, consider
 relaxing the max latency requirement and further scaling the underlying hardware resources.
+
+Custom Batching
+---------------
+
+Currently, adaptive batching is only effective for certain types of parameters including(non-exhaustive) ``numpy.ndarray``, ``pandas.Series`` and framework-specific types such as ``torch.Tensor``.
+Batch parameters of other types are simply collected into a list and passed to the inference function. If your model accepts parameters that are not batchable by default,
+you can achieve adaptive batching by wrapping the inference function with a :ref:`Runner <concepts/runner:Custom Runner>`.
+
+We will demonstrate this with a PyTorch example which accepts a dictionary of ``torch.Tensor`` as input.
+
+.. code-block:: python
+    :caption: `service.py`
+
+    import bentoml
+    import torch
+
+    class MyRunnable(bentoml.Runnable):
+        def __init__(self):
+            super().__init__()
+            # if torch.cuda.device_count():
+            if torch.cuda.is_available():
+                self.device_id = "cuda"
+                # by default, torch.FloatTensor will be used on CPU.
+                torch.set_default_tensor_type("torch.cuda.FloatTensor")
+            else:
+                self.device_id = "cpu"
+            self.model = bentoml.pytorch.load_model(device_id=self.device_id)
+            # We want to turn off dropout and batchnorm when running inference.
+            self.model.train(False)
+
+        @bentoml.Runnable.method(batchable=True)
+        def __call__(self, **batch):
+            # Our model accepts a dictionary as input, but we use ``**``
+            # to deconstruct it to keyword arguments, so they can be batched rightly.
+            # Move the tensors to the target device and pass it to the model as a dict.
+            batch = {k: tensor.to(self.device_id) for k, tensor in batch.items()}
+            with torch.inference_mode():
+                return self.model(batch)
+
+    model = bentoml.pytorch.get("my_pytorch_model")
path: docs/source/guides/batching.rst
text: @user1:
```suggestion
```

@user1:
we don't need this anymore right?

@author:
It may need to be passed to `service`, but for simplicity, yes.

diff_hunk: @@ -114,7 +114,7 @@ We will demonstrate this with a PyTorch example which accepts a dictionary of ``
                 torch.set_default_tensor_type("torch.cuda.FloatTensor")
             else:
                 self.device_id = "cpu"
-            self.model = model.load_model(device_id=self.device_id)
+            self.model = bentoml.pytorch.load_model(device_id=self.device_id)
path: docs/source/guides/batching.rst
text: @user1:
oh sorry, I think you also have to pass in the tag for the Runner, and then pass the tag via `runnable_init_params`.

@author:
For a custom runnable for the usage of the service only, we can hardcode the tag here?

@user2:
Yeah, that would work.