[booster] added the accelerator implementation (#3159)

This commit is contained in:
Frank Lee
2023-03-20 13:59:24 +08:00
committed by GitHub
parent 1ad3a636b1
commit a9b8402d93
4 changed files with 72 additions and 5 deletions

View File

@@ -0,0 +1,13 @@
import pytest
import torch.nn as nn
from torchvision.models import resnet18
from colossalai.booster.accelerator import Accelerator
@pytest.mark.parametrize('device', ['cpu', 'cuda'])
def test_accelerator(device):
acceleartor = Accelerator(device)
model = nn.Linear(8, 8)
model = acceleartor.configure_model(model)
assert next(model.parameters()).device.type == device