提交 c992afa9 编写于 作者: M Megvii Engine Team

test(mge/module): add module reset attribute test

GitOrigin-RevId: 6c9adc4a7022ed453221d7c58d0b4b0033fd00b4
上级 93bfda51
......@@ -625,3 +625,34 @@ def test_repr_module_delete():
del net.softmax
output = net.__repr__()
assert output == ground_truth
def test_repr_module_reset_attr():
class ResetAttrModule(Module):
def __init__(self, flag):
super().__init__()
if flag:
self.a = None
self.a = Linear(3, 5)
else:
self.a = Linear(3, 5)
self.a = None
def forward(self, x):
if self.a:
x = self.a(x)
return x
ground_truth = [
(
"ResetAttrModule(\n"
" (a): Linear(in_features=3, out_features=5, bias=True)\n"
")"
),
("ResetAttrModule()"),
]
m0 = ResetAttrModule(True)
m1 = ResetAttrModule(False)
output = [m0.__repr__(), m1.__repr__()]
assert output == ground_truth
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册