mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-09-05 02:51:59 +00:00
[booster] added the accelerator implementation (#3159)
This commit is contained in:
13
tests/test_booster/test_accelerator.py
Normal file
13
tests/test_booster/test_accelerator.py
Normal 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
|
Reference in New Issue
Block a user