"applications/Chat/coati/vscode:/vscode.git/clone" did not exist on "079bf3cb26a502fc647b1aad15fd14d6266be66c"
test_accelerator.py 387 Bytes
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
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