Fix tutorial example for my dtype tensor subclass (#865)

* Fix tutorial example for my dtyep tensor subclass

Summary:
att

Test Plan:
python tutorials/developer_api_guide/my_dtype_tensor_subclass.py

Reviewers:

Subscribers:

Tasks:

Tags:

* fix function not defined error

Author: jerryzh168

tutorials/developer_api_guide/my_dtype_tensor_subclass.py

@@ -20,7 +20,7 @@
     LayoutType,
     PlainLayoutType,
 )
-from torchao.utils import TorchAOBaseTensor, _register_layout_cls, _get_layout_tensor_constructor
+from torchao.utils import TorchAOBaseTensor
 
 aten = torch.ops.aten
 
@@ -191,12 +191,8 @@ def _apply_fn_to_data(self, fn):
 # LayoutType and Layout Tensor Subclass Registration #
 ######################################################
 
-def register_layout_cls(layout_type_class: type(LayoutType)):
-    return _register_layout_cls(MyDTypeTensor, layout_type_class)
-
-def get_layout_tensor_constructor(layout_type_class: type(LayoutType)):
-    return _get_layout_tensor_constructor(MyDTypeTensor, layout_type_class)
-
+register_layout_cls = MyDTypeTensor.register_layout_cls
+get_layout_tensor_constructor = MyDTypeTensor.get_layout_tensor_constructor
 
 @register_layout_cls(PlainLayoutType)
 class PlainMyDTypeLayout(MyDTypeLayout):
@@ -343,12 +339,12 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
 for _ in range(NUM_WARMUPS):
     m(*example_inputs)
-print("before quantization:", benchmark_model(m, NUM_RUNS, example_inputs[0]))
+print("before quantization:", benchmark_model(m, NUM_RUNS, example_inputs))
 
 compiled = torch.compile(m, mode="max-autotune")
 for _ in range(NUM_WARMUPS):
     compiled(*example_inputs)
-print("after compile:", benchmark_model(compiled, NUM_RUNS, example_inputs[0]))
+print("after compile:", benchmark_model(compiled, NUM_RUNS, example_inputs))
 
 # convert weights to quantized weights
 m.linear.weight = torch.nn.Parameter(
@@ -358,7 +354,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 for _ in range(NUM_WARMUPS):
     m(*example_inputs)
 
-print("after quantization:", benchmark_model(m, NUM_RUNS, example_inputs[0]))
+print("after quantization:", benchmark_model(m, NUM_RUNS, example_inputs))
 
 m = torch.compile(m, mode="max-autotune")
 
@@ -367,4 +363,4 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
 # NOTE: currently there is no speedup because we just dequantize the weight in the _quantized_linear op
 # we plan to add custom op example in the future and that will help us to get speedup
-print("after quantization and compile:", benchmark_model(m, NUM_RUNS, example_inputs[0]))
+print("after quantization and compile:", benchmark_model(m, NUM_RUNS, example_inputs))