mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2026-04-25 01:03:35 +00:00
15 lines
269 B
Python
15 lines
269 B
Python
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
|