mirror of https://github.com/hpcaitech/ColossalAI
[format] applied code formatting on changed files in pull request 4021 (#4022)
Co-authored-by: github-actions <github-actions@github.com>pull/4033/head
parent
160c64c645
commit
a52f62082d
|
@ -1,11 +1,12 @@
|
|||
#!/usr/bin/env python
|
||||
# -*- encoding: utf-8 -*-
|
||||
|
||||
from contextlib import contextmanager
|
||||
|
||||
import torch.nn as nn
|
||||
|
||||
from colossalai.context import ParallelMode
|
||||
from colossalai.core import global_context as gpc
|
||||
from contextlib import contextmanager
|
||||
|
||||
|
||||
class ParallelLayer(nn.Module):
|
||||
|
|
|
@ -1,6 +1,7 @@
|
|||
from colossalai.device.device_mesh import DeviceMesh
|
||||
import torch
|
||||
|
||||
from colossalai.device.device_mesh import DeviceMesh
|
||||
|
||||
|
||||
def test_device_mesh():
|
||||
physical_mesh_id = torch.arange(0, 16).reshape(2, 8)
|
||||
|
|
Loading…
Reference in New Issue