Reverting changes to test_mlp

Author: ign-amanks

models/tt_transformers/tests/test_mlp.py

@@ -9,6 +9,8 @@
 from loguru import logger
 
 import ttnn
+from models.tt_transformers.tests.test_utils import get_ref_model_dype
+from models.tt_transformers.tt.ccl import TT_CCL
 from models.tt_transformers.tt.mlp import MLP
 from models.tt_transformers.tt.model_config import ModelArgs
 from models.utility_functions import comp_allclose, comp_pcc, skip_for_grayskull
@@ -33,11 +35,12 @@
     "batch_size",
     (1,),
 )
+@pytest.mark.parametrize("device_params", [{"fabric_config": True}], indirect=True)
 def test_mlp_inference(seq_len, batch_size, mesh_device, reset_seeds, ensure_gc):
     dtype = ttnn.bfloat8_b
     mode = "decode" if seq_len <= 32 else "prefill"
 
-    model_args = ModelArgs(mesh_device, max_batch_size=batch_size, max_seq_len=128)
+    model_args = ModelArgs(mesh_device, max_batch_size=batch_size, max_seq_len=128, cache_hf=True)
     model_args.n_layers = 1
     state_dict = model_args.load_state_dict()
 
@@ -50,16 +53,21 @@ def test_mlp_inference(seq_len, batch_size, mesh_device, reset_seeds, ensure_gc)
     reference_model = model_args.reference_mlp()
     reference_model.load_state_dict(partial_state_dict)
 
+    tt_ccl = TT_CCL(mesh_device)
     tt_model = MLP(
         mesh_device=mesh_device,
+        tt_ccl=tt_ccl,
         args=model_args,
         state_dict=state_dict,
         weight_cache_path=model_args.weight_cache_path(dtype),
         layer_num=0,
         dtype=dtype,
         model_config=model_args.get_model_config(),
     )
-    torch_input = torch.randn(1, 1, seq_len, model_args.dim)
+
+    torch_input = torch.randn(
+        1, 1, seq_len, model_args.dim, dtype=get_ref_model_dype(reference_model, model_args.model_name)
+    )
     reference_output = reference_model(torch_input)
     tt_input = ttnn.from_torch(
         torch_input,