[Tensor] test parameters() as member function (#896)

tests/test_tensor/test_model.py

@@ -98,6 +98,35 @@ def run_1d_col_tp():
         if i > 5:
             break
 
+# Test the overrided parameters() and named_parameters() member functions
+def test_model_parameters():
+    # build a module with 2 Linear, 4 parameters in total.
+    class Net(torch.nn.Module):
+
+        def __init__(self):
+            super().__init__()
+            self.fcs = torch.nn.Sequential(torch.nn.Linear(2, 3), torch.nn.Linear(3, 2))
+            self.extra_param = torch.nn.Parameter(torch.randn(2))
+
+    with ColoInitContext(device=get_current_device()):
+        model = Net()
+
+    param_cnt = 0
+    for name, p in model.named_parameters():
+        param_cnt += 1
+    assert param_cnt == 5
+
+    param_cnt = 0
+    for name, p in model.named_parameters(recurse=False):
+        param_cnt += 1
+    assert param_cnt == 1
+
+    param_cnt = 0
+    for p in model.fcs[0].parameters(recurse=False):
+        param_cnt += 1
+    assert param_cnt == 2
+
+
 def run_1d_row_tp():
     # A simple net with two stacked nn.Linear
     get_components_func = non_distributed_component_funcs.get_callable('simple_net')
@@ -179,4 +208,5 @@ def test_simple_net(world_size):
 
 
 if __name__ == '__main__':
-    test_simple_net()
+    # test_simple_net()
+    test_model_parameters()