We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[tensor] fix some unittests

[tensor] fix some unittests

[tensor] fix some unittests

[booster] implement accelerator
# Overview


`Accelerator` is an abstraction for the hardware devices used for training. It can be CPU, GPU, TPU, etc. With this abstraction, we can provide room for future extension to more hardware which is compatible with PyTorch.

Wanna track the development progress? Take a look at

proposal: https://github.com/hpcaitech/ColossalAI/discussions/3046
project kanban: https://github.com/orgs/hpcaitech/projects/19

# Goal

The `Accelerator` can move models to the target device. A sample definition is given below and subject to future changes.

```python

class Accelerator:

        def __init__(self, device):
              self.device = device

        def setup_model(self, model) -> torch.nn.Module
```



--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `colossalai/booster/accelerator.py`
Content:
```
1 import torch
2 import torch.nn as nn
3 
4 __all__ = ['Accelerator']
5 
6 
7 class Accelerator:
8 
9     def __init__(self, device: torch.device):
10         self.device = device
11 
12     def setup_model(self, model: nn.Module) -> nn.Module:
13         # TODO: implement this method
14         pass
15
```

Path: `colossalai/booster/booster.py`
Content:
```
1 from contextlib import contextmanager
2 from typing import Callable, Iterable, Iterator, List, Optional, Tuple, Union
3 
4 import torch
5 import torch.nn as nn
6 from torch import Tensor
7 from torch.optim import Optimizer
8 from torch.optim.lr_scheduler import _LRScheduler as LRScheduler
9 from torch.utils.data import DataLoader
10 
11 from .mixed_precision import MixedPrecision, mixed_precision_factory
12 from .plugin import Plugin
13 
14 __all__ = ['Booster']
15 
16 
17 class Booster:
18     """
19     Booster is a high-level API for training neural networks. It provides a unified interface for
20     training with different precisio, accelerator, and plugin.
21 
22     Examples:
23         >>> colossalai.launch(...)
24         >>> plugin = GeminiPlugin(stage=3, ...)
25         >>> booster = Booster(precision='fp16', plugin=plugin)
26         >>>
27         >>> model = GPT2()
28         >>> optimizer = Adam(model.parameters())
29         >>> dataloader = Dataloader(Dataset)
30         >>> lr_scheduler = LinearWarmupScheduler()
31         >>> criterion = GPTLMLoss()
32         >>>
33         >>> model, optimizer, lr_scheduler, dataloader = booster.boost(model, optimizer, lr_scheduler, dataloader)
34         >>>
35         >>> for epoch in range(max_epochs):
36         >>>     for input_ids, attention_mask in dataloader:
37         >>>         outputs = model(input_ids, attention_mask)
38         >>>         loss = criterion(outputs.logits, input_ids)
39         >>>         booster.backward(loss, optimizer)
40         >>>         optimizer.step()
41         >>>         lr_scheduler.step()
42         >>>         optimizer.zero_grad()
43 
44 
45     Args:
46         device (str or torch.device): The device to run the training. Default: 'cuda'.
47         mixed_precision (str or MixedPrecision): The mixed precision to run the training. Default: None.
48                                 If the argument is a string, it can be 'fp16', 'fp16_apex', 'bf16', or 'fp8'.
49                                 'fp16' would use PyTorch AMP while `fp16_apex` would use Nvidia Apex.
50         plugin (Plugin): The plugin to run the training. Default: None.
51     """
52 
53     def __init__(self,
54                  device: Union[str, torch.device] = 'cuda',
55                  mixed_precision: Union[MixedPrecision, str] = None,
56                  plugin: Optional[Plugin] = None) -> None:
57         # validate and set precision
58         if isinstance(MixedPrecision, str):
59             # the user will take the default arguments for amp training
60             self.mixed_precision = mixed_precision_factory(mixed_precision)
61         elif isinstance(mixed_precision, MixedPrecision):
62             # the user can customize the arguments by passing the precision object
63             self.mixed_precision = mixed_precision
64         else:
65             raise ValueError(
66                 f'Expected the argument mixed_precision to be a string or an instance of Precision, but got {type(mixed_precision)}.'
67             )
68 
69     def boost(self, model: nn.Module, optimizer: Optimizer, criterion: Callable, lr_scheduler: LRScheduler,
70               dataloader: DataLoader) -> List[Union[nn.Module, Optimizer, LRScheduler, DataLoader]]:
71         """
72         Boost the model, optimizer, criterion, lr_scheduler, and dataloader.
73 
74         Args:
75             model (nn.Module): The model to be boosted.
76             optimizer (Optimizer): The optimizer to be boosted.
77             criterion (Callable): The criterion to be boosted.
78             lr_scheduler (LRScheduler): The lr_scheduler to be boosted.
79             dataloader (DataLoader): The dataloader to be boosted.
80         """
81         # TODO(FrankLeeeee): consider multi-model and multi-optimizer case
82         # TODO(lsg): Add plugin control logic
83         # e.g.
84         # if self.plugin is not None and self.plugin.control_boost:
85         #    ...
86         # transform model for mixed precision
87         model, optimizer, criterion = self.mixed_precision.configure(model, optimizer, criterion)
88         return model, optimizer, criterion, lr_scheduler, dataloader
89 
90     def backward(self, loss: torch.Tensor, optimizer: Optimizer) -> None:
91         # TODO: implement this method with plugin
92         optimizer.backward(loss)
93 
94     def execute_pipeline(self,
95                          data_iter: Iterator,
96                          model: nn.Module,
97                          criterion: Callable[[torch.Tensor], torch.Tensor],
98                          optimizer: Optimizer,
99                          return_loss: bool = True,
100                          return_outputs: bool = False) -> Tuple[Optional[torch.Tensor], ...]:
101         # TODO: implement this method
102         # run pipeline forward backward pass
103         # return loss or outputs if needed
104         pass
105 
106     def no_sync(self, model: nn.Module) -> contextmanager:
107         # TODO: implement this method
108         pass
109 
110     def save(self,
111              obj: Union[nn.Module, Optimizer, LRScheduler],
112              path_like: str,
113              plan: str = 'torch',
114              **kwargs) -> None:
115         # TODO: implement this method
116         pass
117 
118     def load(self,
119              obj: Union[nn.Module, Optimizer, LRScheduler],
120              path_like: str,
121              plan: str = 'torch',
122              **kwargs) -> None:
123         # TODO: implement this method
124         pass
125
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/colossalai/booster/accelerator.py b/colossalai/booster/accelerator.py
--- a/colossalai/booster/accelerator.py
+++ b/colossalai/booster/accelerator.py
@@ -3,12 +3,52 @@
 
 __all__ = ['Accelerator']
 
+_supported_devices = [
+    'cpu',
+    'cuda',
+
+    # To be supported
+    # 'xpu',
+    # 'npu',
+    # 'tpu',
+]
+
 
 class Accelerator:
+    """
+    Accelerator is an abstraction for the hardware device that is used to run the model.
+
+    Args:
+        device (str): The device to be used. Currently only support 'cpu' and 'gpu'.
+    """
 
-    def __init__(self, device: torch.device):
+    def __init__(self, device: str):
         self.device = device
 
-    def setup_model(self, model: nn.Module) -> nn.Module:
-        # TODO: implement this method
-        pass
+        assert self.device in _supported_devices, f"Device {self.device} is not supported yet, supported devices include {_supported_devices}"
+
+    def bind(self):
+        """
+        Set the default device for the current process.
+        """
+        if self.device == 'cpu':
+            pass
+        elif self.device == 'cuda':
+            # TODO(FrankLeeeee): use global environment to check if it is a dist job
+            # if is_distributed:
+            #     local_rank = EnvTable().get_local_rank()
+            #     torch.cuda.set_device(torch.device(f'cuda:{local_rank}'))
+            torch.cuda.set_device(torch.device('cuda'))
+            pass
+        else:
+            raise ValueError(f"Device {self.device} is not supported yet")
+
+    def configure_model(self, model: nn.Module) -> nn.Module:
+        """
+        Move the model to the device.
+
+        Args:
+            model (nn.Module): The model to be moved.
+        """
+        model = model.to(torch.device(self.device))
+        return model
diff --git a/colossalai/booster/booster.py b/colossalai/booster/booster.py
--- a/colossalai/booster/booster.py
+++ b/colossalai/booster/booster.py
@@ -8,6 +8,7 @@
 from torch.optim.lr_scheduler import _LRScheduler as LRScheduler
 from torch.utils.data import DataLoader
 
+from .accelerator import Accelerator
 from .mixed_precision import MixedPrecision, mixed_precision_factory
 from .plugin import Plugin
 
@@ -51,9 +52,16 @@
     """
 
     def __init__(self,
-                 device: Union[str, torch.device] = 'cuda',
+                 device: str = 'cuda',
                  mixed_precision: Union[MixedPrecision, str] = None,
                  plugin: Optional[Plugin] = None) -> None:
+        # TODO(FrankLeeeee): add plugin control logic
+        # if self.plugin is not None and self.plugin.control_accelerator:
+        #     ...
+        # create acclerator
+        self.acceleartor = Accelerator(device)
+        self.acceleartor.set_default_device()
+
         # validate and set precision
         if isinstance(MixedPrecision, str):
             # the user will take the default arguments for amp training
@@ -78,6 +86,11 @@
             lr_scheduler (LRScheduler): The lr_scheduler to be boosted.
             dataloader (DataLoader): The dataloader to be boosted.
         """
+        # TODO(FrankLeeeee): add plugin control logic
+        # if self.plugin is not None and self.plugin.control_accelerator:
+        #     ...
+        model = self.acceleartor.configure_model(model)
+
         # TODO(FrankLeeeee): consider multi-model and multi-optimizer case
         # TODO(lsg): Add plugin control logic
         # e.g.

{"golden_diff": "diff --git a/colossalai/booster/accelerator.py b/colossalai/booster/accelerator.py\n--- a/colossalai/booster/accelerator.py\n+++ b/colossalai/booster/accelerator.py\n@@ -3,12 +3,52 @@\n \n __all__ = ['Accelerator']\n \n+_supported_devices = [\n+    'cpu',\n+    'cuda',\n+\n+    # To be supported\n+    # 'xpu',\n+    # 'npu',\n+    # 'tpu',\n+]\n+\n \n class Accelerator:\n+    \"\"\"\n+    Accelerator is an abstraction for the hardware device that is used to run the model.\n+\n+    Args:\n+        device (str): The device to be used. Currently only support 'cpu' and 'gpu'.\n+    \"\"\"\n \n-    def __init__(self, device: torch.device):\n+    def __init__(self, device: str):\n         self.device = device\n \n-    def setup_model(self, model: nn.Module) -> nn.Module:\n-        # TODO: implement this method\n-        pass\n+        assert self.device in _supported_devices, f\"Device {self.device} is not supported yet, supported devices include {_supported_devices}\"\n+\n+    def bind(self):\n+        \"\"\"\n+        Set the default device for the current process.\n+        \"\"\"\n+        if self.device == 'cpu':\n+            pass\n+        elif self.device == 'cuda':\n+            # TODO(FrankLeeeee): use global environment to check if it is a dist job\n+            # if is_distributed:\n+            #     local_rank = EnvTable().get_local_rank()\n+            #     torch.cuda.set_device(torch.device(f'cuda:{local_rank}'))\n+            torch.cuda.set_device(torch.device('cuda'))\n+            pass\n+        else:\n+            raise ValueError(f\"Device {self.device} is not supported yet\")\n+\n+    def configure_model(self, model: nn.Module) -> nn.Module:\n+        \"\"\"\n+        Move the model to the device.\n+\n+        Args:\n+            model (nn.Module): The model to be moved.\n+        \"\"\"\n+        model = model.to(torch.device(self.device))\n+        return model\ndiff --git a/colossalai/booster/booster.py b/colossalai/booster/booster.py\n--- a/colossalai/booster/booster.py\n+++ b/colossalai/booster/booster.py\n@@ -8,6 +8,7 @@\n from torch.optim.lr_scheduler import _LRScheduler as LRScheduler\n from torch.utils.data import DataLoader\n \n+from .accelerator import Accelerator\n from .mixed_precision import MixedPrecision, mixed_precision_factory\n from .plugin import Plugin\n \n@@ -51,9 +52,16 @@\n     \"\"\"\n \n     def __init__(self,\n-                 device: Union[str, torch.device] = 'cuda',\n+                 device: str = 'cuda',\n                  mixed_precision: Union[MixedPrecision, str] = None,\n                  plugin: Optional[Plugin] = None) -> None:\n+        # TODO(FrankLeeeee): add plugin control logic\n+        # if self.plugin is not None and self.plugin.control_accelerator:\n+        #     ...\n+        # create acclerator\n+        self.acceleartor = Accelerator(device)\n+        self.acceleartor.set_default_device()\n+\n         # validate and set precision\n         if isinstance(MixedPrecision, str):\n             # the user will take the default arguments for amp training\n@@ -78,6 +86,11 @@\n             lr_scheduler (LRScheduler): The lr_scheduler to be boosted.\n             dataloader (DataLoader): The dataloader to be boosted.\n         \"\"\"\n+        # TODO(FrankLeeeee): add plugin control logic\n+        # if self.plugin is not None and self.plugin.control_accelerator:\n+        #     ...\n+        model = self.acceleartor.configure_model(model)\n+\n         # TODO(FrankLeeeee): consider multi-model and multi-optimizer case\n         # TODO(lsg): Add plugin control logic\n         # e.g.\n", "issue": "[tensor] fix some unittests\n\n[tensor] fix some unittests\n\n[tensor] fix some unittests\n\n[booster] implement accelerator\n# Overview\n\n\n`Accelerator` is an abstraction for the hardware devices used for training. It can be CPU, GPU, TPU, etc. With this abstraction, we can provide room for future extension to more hardware which is compatible with PyTorch.\n\nWanna track the development progress? Take a look at\n\nproposal: https://github.com/hpcaitech/ColossalAI/discussions/3046\nproject kanban: https://github.com/orgs/hpcaitech/projects/19\n\n# Goal\n\nThe `Accelerator` can move models to the target device. A sample definition is given below and subject to future changes.\n\n```python\n\nclass Accelerator:\n\n        def __init__(self, device):\n              self.device = device\n\n        def setup_model(self, model) -> torch.nn.Module\n```\n\n\n", "before_files": [{"content": "import torch\nimport torch.nn as nn\n\n__all__ = ['Accelerator']\n\n\nclass Accelerator:\n\n    def __init__(self, device: torch.device):\n        self.device = device\n\n    def setup_model(self, model: nn.Module) -> nn.Module:\n        # TODO: implement this method\n        pass\n", "path": "colossalai/booster/accelerator.py"}, {"content": "from contextlib import contextmanager\nfrom typing import Callable, Iterable, Iterator, List, Optional, Tuple, Union\n\nimport torch\nimport torch.nn as nn\nfrom torch import Tensor\nfrom torch.optim import Optimizer\nfrom torch.optim.lr_scheduler import _LRScheduler as LRScheduler\nfrom torch.utils.data import DataLoader\n\nfrom .mixed_precision import MixedPrecision, mixed_precision_factory\nfrom .plugin import Plugin\n\n__all__ = ['Booster']\n\n\nclass Booster:\n    \"\"\"\n    Booster is a high-level API for training neural networks. It provides a unified interface for\n    training with different precisio, accelerator, and plugin.\n\n    Examples:\n        >>> colossalai.launch(...)\n        >>> plugin = GeminiPlugin(stage=3, ...)\n        >>> booster = Booster(precision='fp16', plugin=plugin)\n        >>>\n        >>> model = GPT2()\n        >>> optimizer = Adam(model.parameters())\n        >>> dataloader = Dataloader(Dataset)\n        >>> lr_scheduler = LinearWarmupScheduler()\n        >>> criterion = GPTLMLoss()\n        >>>\n        >>> model, optimizer, lr_scheduler, dataloader = booster.boost(model, optimizer, lr_scheduler, dataloader)\n        >>>\n        >>> for epoch in range(max_epochs):\n        >>>     for input_ids, attention_mask in dataloader:\n        >>>         outputs = model(input_ids, attention_mask)\n        >>>         loss = criterion(outputs.logits, input_ids)\n        >>>         booster.backward(loss, optimizer)\n        >>>         optimizer.step()\n        >>>         lr_scheduler.step()\n        >>>         optimizer.zero_grad()\n\n\n    Args:\n        device (str or torch.device): The device to run the training. Default: 'cuda'.\n        mixed_precision (str or MixedPrecision): The mixed precision to run the training. Default: None.\n                                If the argument is a string, it can be 'fp16', 'fp16_apex', 'bf16', or 'fp8'.\n                                'fp16' would use PyTorch AMP while `fp16_apex` would use Nvidia Apex.\n        plugin (Plugin): The plugin to run the training. Default: None.\n    \"\"\"\n\n    def __init__(self,\n                 device: Union[str, torch.device] = 'cuda',\n                 mixed_precision: Union[MixedPrecision, str] = None,\n                 plugin: Optional[Plugin] = None) -> None:\n        # validate and set precision\n        if isinstance(MixedPrecision, str):\n            # the user will take the default arguments for amp training\n            self.mixed_precision = mixed_precision_factory(mixed_precision)\n        elif isinstance(mixed_precision, MixedPrecision):\n            # the user can customize the arguments by passing the precision object\n            self.mixed_precision = mixed_precision\n        else:\n            raise ValueError(\n                f'Expected the argument mixed_precision to be a string or an instance of Precision, but got {type(mixed_precision)}.'\n            )\n\n    def boost(self, model: nn.Module, optimizer: Optimizer, criterion: Callable, lr_scheduler: LRScheduler,\n              dataloader: DataLoader) -> List[Union[nn.Module, Optimizer, LRScheduler, DataLoader]]:\n        \"\"\"\n        Boost the model, optimizer, criterion, lr_scheduler, and dataloader.\n\n        Args:\n            model (nn.Module): The model to be boosted.\n            optimizer (Optimizer): The optimizer to be boosted.\n            criterion (Callable): The criterion to be boosted.\n            lr_scheduler (LRScheduler): The lr_scheduler to be boosted.\n            dataloader (DataLoader): The dataloader to be boosted.\n        \"\"\"\n        # TODO(FrankLeeeee): consider multi-model and multi-optimizer case\n        # TODO(lsg): Add plugin control logic\n        # e.g.\n        # if self.plugin is not None and self.plugin.control_boost:\n        #    ...\n        # transform model for mixed precision\n        model, optimizer, criterion = self.mixed_precision.configure(model, optimizer, criterion)\n        return model, optimizer, criterion, lr_scheduler, dataloader\n\n    def backward(self, loss: torch.Tensor, optimizer: Optimizer) -> None:\n        # TODO: implement this method with plugin\n        optimizer.backward(loss)\n\n    def execute_pipeline(self,\n                         data_iter: Iterator,\n                         model: nn.Module,\n                         criterion: Callable[[torch.Tensor], torch.Tensor],\n                         optimizer: Optimizer,\n                         return_loss: bool = True,\n                         return_outputs: bool = False) -> Tuple[Optional[torch.Tensor], ...]:\n        # TODO: implement this method\n        # run pipeline forward backward pass\n        # return loss or outputs if needed\n        pass\n\n    def no_sync(self, model: nn.Module) -> contextmanager:\n        # TODO: implement this method\n        pass\n\n    def save(self,\n             obj: Union[nn.Module, Optimizer, LRScheduler],\n             path_like: str,\n             plan: str = 'torch',\n             **kwargs) -> None:\n        # TODO: implement this method\n        pass\n\n    def load(self,\n             obj: Union[nn.Module, Optimizer, LRScheduler],\n             path_like: str,\n             plan: str = 'torch',\n             **kwargs) -> None:\n        # TODO: implement this method\n        pass\n", "path": "colossalai/booster/booster.py"}], "after_files": [{"content": "import torch\nimport torch.nn as nn\n\n__all__ = ['Accelerator']\n\n_supported_devices = [\n    'cpu',\n    'cuda',\n\n    # To be supported\n    # 'xpu',\n    # 'npu',\n    # 'tpu',\n]\n\n\nclass Accelerator:\n    \"\"\"\n    Accelerator is an abstraction for the hardware device that is used to run the model.\n\n    Args:\n        device (str): The device to be used. Currently only support 'cpu' and 'gpu'.\n    \"\"\"\n\n    def __init__(self, device: str):\n        self.device = device\n\n        assert self.device in _supported_devices, f\"Device {self.device} is not supported yet, supported devices include {_supported_devices}\"\n\n    def bind(self):\n        \"\"\"\n        Set the default device for the current process.\n        \"\"\"\n        if self.device == 'cpu':\n            pass\n        elif self.device == 'cuda':\n            # TODO(FrankLeeeee): use global environment to check if it is a dist job\n            # if is_distributed:\n            #     local_rank = EnvTable().get_local_rank()\n            #     torch.cuda.set_device(torch.device(f'cuda:{local_rank}'))\n            torch.cuda.set_device(torch.device('cuda'))\n            pass\n        else:\n            raise ValueError(f\"Device {self.device} is not supported yet\")\n\n    def configure_model(self, model: nn.Module) -> nn.Module:\n        \"\"\"\n        Move the model to the device.\n\n        Args:\n            model (nn.Module): The model to be moved.\n        \"\"\"\n        model = model.to(torch.device(self.device))\n        return model\n", "path": "colossalai/booster/accelerator.py"}, {"content": "from contextlib import contextmanager\nfrom typing import Callable, Iterable, Iterator, List, Optional, Tuple, Union\n\nimport torch\nimport torch.nn as nn\nfrom torch import Tensor\nfrom torch.optim import Optimizer\nfrom torch.optim.lr_scheduler import _LRScheduler as LRScheduler\nfrom torch.utils.data import DataLoader\n\nfrom .accelerator import Accelerator\nfrom .mixed_precision import MixedPrecision, mixed_precision_factory\nfrom .plugin import Plugin\n\n__all__ = ['Booster']\n\n\nclass Booster:\n    \"\"\"\n    Booster is a high-level API for training neural networks. It provides a unified interface for\n    training with different precisio, accelerator, and plugin.\n\n    Examples:\n        >>> colossalai.launch(...)\n        >>> plugin = GeminiPlugin(stage=3, ...)\n        >>> booster = Booster(precision='fp16', plugin=plugin)\n        >>>\n        >>> model = GPT2()\n        >>> optimizer = Adam(model.parameters())\n        >>> dataloader = Dataloader(Dataset)\n        >>> lr_scheduler = LinearWarmupScheduler()\n        >>> criterion = GPTLMLoss()\n        >>>\n        >>> model, optimizer, lr_scheduler, dataloader = booster.boost(model, optimizer, lr_scheduler, dataloader)\n        >>>\n        >>> for epoch in range(max_epochs):\n        >>>     for input_ids, attention_mask in dataloader:\n        >>>         outputs = model(input_ids, attention_mask)\n        >>>         loss = criterion(outputs.logits, input_ids)\n        >>>         booster.backward(loss, optimizer)\n        >>>         optimizer.step()\n        >>>         lr_scheduler.step()\n        >>>         optimizer.zero_grad()\n\n\n    Args:\n        device (str or torch.device): The device to run the training. Default: 'cuda'.\n        mixed_precision (str or MixedPrecision): The mixed precision to run the training. Default: None.\n                                If the argument is a string, it can be 'fp16', 'fp16_apex', 'bf16', or 'fp8'.\n                                'fp16' would use PyTorch AMP while `fp16_apex` would use Nvidia Apex.\n        plugin (Plugin): The plugin to run the training. Default: None.\n    \"\"\"\n\n    def __init__(self,\n                 device: str = 'cuda',\n                 mixed_precision: Union[MixedPrecision, str] = None,\n                 plugin: Optional[Plugin] = None) -> None:\n        # TODO(FrankLeeeee): add plugin control logic\n        # if self.plugin is not None and self.plugin.control_accelerator:\n        #     ...\n        # create acclerator\n        self.acceleartor = Accelerator(device)\n        self.acceleartor.set_default_device()\n\n        # validate and set precision\n        if isinstance(MixedPrecision, str):\n            # the user will take the default arguments for amp training\n            self.mixed_precision = mixed_precision_factory(mixed_precision)\n        elif isinstance(mixed_precision, MixedPrecision):\n            # the user can customize the arguments by passing the precision object\n            self.mixed_precision = mixed_precision\n        else:\n            raise ValueError(\n                f'Expected the argument mixed_precision to be a string or an instance of Precision, but got {type(mixed_precision)}.'\n            )\n\n    def boost(self, model: nn.Module, optimizer: Optimizer, criterion: Callable, lr_scheduler: LRScheduler,\n              dataloader: DataLoader) -> List[Union[nn.Module, Optimizer, LRScheduler, DataLoader]]:\n        \"\"\"\n        Boost the model, optimizer, criterion, lr_scheduler, and dataloader.\n\n        Args:\n            model (nn.Module): The model to be boosted.\n            optimizer (Optimizer): The optimizer to be boosted.\n            criterion (Callable): The criterion to be boosted.\n            lr_scheduler (LRScheduler): The lr_scheduler to be boosted.\n            dataloader (DataLoader): The dataloader to be boosted.\n        \"\"\"\n        # TODO(FrankLeeeee): add plugin control logic\n        # if self.plugin is not None and self.plugin.control_accelerator:\n        #     ...\n        model = self.acceleartor.configure_model(model)\n\n        # TODO(FrankLeeeee): consider multi-model and multi-optimizer case\n        # TODO(lsg): Add plugin control logic\n        # e.g.\n        # if self.plugin is not None and self.plugin.control_boost:\n        #    ...\n        # transform model for mixed precision\n        model, optimizer, criterion = self.mixed_precision.configure(model, optimizer, criterion)\n        return model, optimizer, criterion, lr_scheduler, dataloader\n\n    def backward(self, loss: torch.Tensor, optimizer: Optimizer) -> None:\n        # TODO: implement this method with plugin\n        optimizer.backward(loss)\n\n    def execute_pipeline(self,\n                         data_iter: Iterator,\n                         model: nn.Module,\n                         criterion: Callable[[torch.Tensor], torch.Tensor],\n                         optimizer: Optimizer,\n                         return_loss: bool = True,\n                         return_outputs: bool = False) -> Tuple[Optional[torch.Tensor], ...]:\n        # TODO: implement this method\n        # run pipeline forward backward pass\n        # return loss or outputs if needed\n        pass\n\n    def no_sync(self, model: nn.Module) -> contextmanager:\n        # TODO: implement this method\n        pass\n\n    def save(self,\n             obj: Union[nn.Module, Optimizer, LRScheduler],\n             path_like: str,\n             plan: str = 'torch',\n             **kwargs) -> None:\n        # TODO: implement this method\n        pass\n\n    def load(self,\n             obj: Union[nn.Module, Optimizer, LRScheduler],\n             path_like: str,\n             plan: str = 'torch',\n             **kwargs) -> None:\n        # TODO: implement this method\n        pass\n", "path": "colossalai/booster/booster.py"}]}