github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-11 13:59:08 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Tensor ] Add 1Drow weight reshard by spec (#854)

Browse Source
This commit is contained in:
Ziyue Jiang
2022-04-24 18:30:20 +08:00
committed by GitHub
parent d7e0303d1e
commit bcc8655021
5 changed files with 41 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
tests/components_to_test/simple_net.py
View File

@@ -1,10 +1,11 @@
from zmq import device
import torch
import torch.nn as nn
import torch.nn.functional as F
from colossalai.nn import CheckpointModule
from .utils.dummy_data_generator import DummyDataGenerator
from .registry import non_distributed_component_funcs
from colossalai.utils.cuda import get_current_device
class SimpleNet(CheckpointModule):
"""
@@ -25,8 +26,8 @@ class SimpleNet(CheckpointModule):
class DummyDataLoader(DummyDataGenerator):
def generate(self):
data = torch.rand(16, 4)
label = torch.randint(low=0, high=2, size=(16,))
data = torch.rand(16, 4, device=get_current_device())
label = torch.randint(low=0, high=2, size=(16,), device=get_current_device())
return data, label