Update babyllama generation

Differential Revision: D83066067

Pull Request resolved: #14519

Author: mcremon-meta

examples/cadence/models/babyllama.py

@@ -14,8 +14,10 @@
 
 from executorch.backends.cadence.aot.export_example import export_and_run_model
 
-from executorch.examples.models.llama.llama_transformer import ModelArgs, Transformer
-
+from executorch.examples.models.llama.llama_transformer import (
+    construct_transformer,
+    ModelArgs,
+)
 
 FORMAT = "[%(levelname)s %(asctime)s %(filename)s:%(lineno)s] %(message)s"
 logging.basicConfig(level=logging.INFO, format=FORMAT)
@@ -32,7 +34,7 @@ def main() -> None:
     )
     seq = 64
     b = 1
-    model = Transformer(args)
+    model = construct_transformer(args)
     example_inputs = (torch.randint(0, 10, [b, seq], dtype=torch.int64),)
 
     export_and_run_model(model, example_inputs)