fix example nn_module

python/tvm/relax/testing/nn.py

@@ -119,14 +119,14 @@ def _unpack_params(value: object) -> List[relax.Var]:
 
 def init_params(mod: tvm.IRModule) -> List[tvm.nd.array]:
     """Utility function to initialize model's parameters."""
-    shape_dict = {v.name_hint: v.shape_ for v in mod["main"].params}
+    sinfo_dict = {v.name_hint: v.struct_info for v in mod["main"].params}
     params = []
-    for k, v in shape_dict.items():
+    for k, v in sinfo_dict.items():
         if k.startswith("data"):
             continue
-        if isinstance(v, relax.ShapeExpr):
+        if isinstance(v, relax.TensorStructInfo) and isinstance(v.shape, relax.ShapeExpr):
             shape = []
-            for i in v:
+            for i in v.shape:
                 if isinstance(i, tir.IntImm):
                     shape.append(int(i))
                 else: