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

[ci] cache cuda extension (#860)

Browse Source
This commit is contained in:
Frank Lee
2022-04-25 10:03:47 +08:00
committed by GitHub
parent c6930d8ddf
commit 1258af71cc
2 changed files with 5 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
tests/components_to_test/simple_net.py
View File

@@ -1,12 +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):
"""
In this no-leaf module, it has subordinate nn.modules and a nn.Parameter.