github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-04 02:26:51 +00:00
Code Issues Packages Projects Releases Wiki Activity

[misc] update pre-commit and run all files (#4752)

Browse Source 
* [misc] update pre-commit

* [misc] run pre-commit

* [misc] remove useless configuration files

* [misc] ignore cuda for clang-format
This commit is contained in:
Hongxin Liu
2023-09-19 14:20:26 +08:00
committed by GitHub
parent 3c6b831c26
commit 079bf3cb26
1268 changed files with 50037 additions and 38444 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
tests/test_fx/test_coloproxy.py
View File

@@ -1,4 +1,3 @@
import pytest
import torch
import torch.nn as nn
from torch.fx import GraphModule
@@ -9,7 +8,6 @@ from colossalai.testing import clear_cache_before_run
class Conv1D(nn.Module):
def __init__(self, nf, nx):
super().__init__()
self.nf = nf
@@ -27,10 +25,9 @@ class Conv1D(nn.Module):
@clear_cache_before_run()
def test_coloproxy():
tracer = ColoTracer()
model = Conv1D(3, 3)
input_sample = {'x': torch.rand(3, 3).to('meta')}
input_sample = {"x": torch.rand(3, 3).to("meta")}
graph = tracer.trace(root=model, meta_args=input_sample)
gm = GraphModule(model, graph, model.__class__.__name__)
@@ -38,7 +35,7 @@ def test_coloproxy():
node = list(gm.graph.nodes)[0]
proxy = ColoProxy(node=node, tracer=tracer)
proxy.meta_data = torch.empty(4, 2, device='meta')
proxy.meta_data = torch.empty(4, 2, device="meta")
assert len(proxy) == 4
assert proxy.shape[0] == 4 and proxy.shape[1] == 2
@@ -47,5 +44,5 @@ def test_coloproxy():
assert proxy.size(0) == 4
if __name__ == '__main__':
if __name__ == "__main__":
test_coloproxy()