diff --git a/tests/nn/tensor_parallel/conftest.py b/tests/nn/tensor_parallel/conftest.py
index 820719d..2b03150 100644
--- a/tests/nn/tensor_parallel/conftest.py
+++ b/tests/nn/tensor_parallel/conftest.py
@@ -1,9 +1,9 @@
 import pytest
-from transformers import AutoModel
+from transformers import AutoModelForCausalLM
 
 MODEL_NAME = "bigscience/bloom-560m"
 
 
 @pytest.fixture(scope="session")
 def model():
-    return AutoModel.from_pretrained(MODEL_NAME)
+    return AutoModelForCausalLM.from_pretrained(MODEL_NAME)
diff --git a/tests/nn/tensor_parallel/test_parallelize.py b/tests/nn/tensor_parallel/test_parallelize.py
index 9b4b894..ca056f1 100644
--- a/tests/nn/tensor_parallel/test_parallelize.py
+++ b/tests/nn/tensor_parallel/test_parallelize.py
@@ -85,10 +85,10 @@ def run_parallelize_linear(
 
 
 @pytest.mark.parametrize("tensor_parallel_size, MODULE_NAME, get_module", [
-    (1, "transformer.h.0.mlp.dense_h_to_4h", lambda model: model.h[0].mlp.dense_h_to_4h),
-    (2, "transformer.h.0.mlp.dense_h_to_4h", lambda model: model.h[0].mlp.dense_h_to_4h),
-    (1, "transformer.h.0.mlp.dense_4h_to_h", lambda model: model.h[0].mlp.dense_4h_to_h),
-    (2, "transformer.h.0.mlp.dense_4h_to_h", lambda model: model.h[0].mlp.dense_4h_to_h),
+    (1, "transformer.h.0.mlp.dense_h_to_4h", lambda model: model.transformer.h[0].mlp.dense_h_to_4h),
+    (2, "transformer.h.0.mlp.dense_h_to_4h", lambda model: model.transformer.h[0].mlp.dense_h_to_4h),
+    (1, "transformer.h.0.mlp.dense_4h_to_h", lambda model: model.transformer.h[0].mlp.dense_4h_to_h),
+    (2, "transformer.h.0.mlp.dense_4h_to_h", lambda model: model.transformer.h[0].mlp.dense_4h_to_h),
 ])
 def test_parallelize_linear(model, tensor_parallel_size, MODULE_NAME, get_module):
     PIPELINE_PARALLEL_SIZE = 1