Reduce memory requirement on export_llama tests with no params (pytorch#11175)

Differential Revision: D75498713

Pull Request resolved: pytorch#11339

Author: jackzhxng

examples/models/llama/model_args.py

@@ -5,7 +5,7 @@
 @dataclass
 class ModelArgs:
     dim: int = 4096
-    n_layers: int = 8
+    n_layers: int = 1
     n_heads: int = 32
     n_kv_heads: Optional[int] = None
     vocab_size: int = 512  # Arbitrary value, should be defined later by tokenizer.

examples/models/llama/tests/test_export_llama_lib.py

@@ -19,7 +19,6 @@
 
 
 class ExportLlamaLibTest(unittest.TestCase):
-    @unittest.skip("Keeps failing on trunk, temporarily skip")
     def test_has_expected_ops_and_op_counts(self):
         """
         Checks the presence of unwanted expensive ops.