mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-09-11 13:59:08 +00:00
[booster] init module structure and definition (#3056)
This commit is contained in:
14
colossalai/booster/accelerator.py
Normal file
14
colossalai/booster/accelerator.py
Normal file
@@ -0,0 +1,14 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
__all__ = ['Accelerator']
|
||||
|
||||
|
||||
class Accelerator:
|
||||
|
||||
def __init__(self, device: torch.device):
|
||||
self.device = device
|
||||
|
||||
def setup_model(self, model: nn.Module) -> nn.Module:
|
||||
# TODO: implement this method
|
||||
pass
|
Reference in New Issue
Block a user