diff --git a/projects/mock_transformers/dist_infer_opt.py b/projects/mock_transformers/dist_infer_opt.py
index 6f03af86d..4bd8728b4 100644
--- a/projects/mock_transformers/dist_infer_opt.py
+++ b/projects/mock_transformers/dist_infer_opt.py
@@ -81,13 +81,20 @@ def __init__(self, *args, **kwargs):
     )
     dist.setup_dist_util(parallel_config)
 
-    # initial and load model
-    model = AutoModelForCausalLM.from_pretrained("facebook/opt-2.7b", torch_dtype=flow.float16)
+    placement_sbp_dict = dict(
+        placement=flow.env.all_device_placement("cuda"),
+        sbp=flow.sbp.broadcast,
+    )
+
+    with global_mode(True, **placement_sbp_dict):
+        # initial and load model
+        model = AutoModelForCausalLM.from_pretrained("facebook/opt-125m", torch_dtype=flow.float16)
+
     # set model to cuda
     dist.set_device_type("cuda")
     model._apply(dist.convert_to_distributed_default_setting)
     # initial tokenizer
-    tokenizer = AutoTokenizer.from_pretrained("facebook/opt-2.7b", use_fast=False)
+    tokenizer = AutoTokenizer.from_pretrained("facebook/opt-125m", use_fast=False)
 
     # get input_ids
     prompt = "Hello, I'm am conscious and"
@@ -98,11 +105,6 @@ def __init__(self, *args, **kwargs):
         placement=dist.get_layer_placement(0),
     )
 
-    # generate id
-    placement_sbp_dict = dict(
-        placement=flow.env.all_device_placement("cuda"),
-        sbp=flow.sbp.broadcast,
-    )
     with global_mode(True, **placement_sbp_dict):
         generated_ids = model.generate(input_ids, max_length=30)
     out_put_ids = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
diff --git a/projects/mock_transformers/init_env.py b/projects/mock_transformers/init_env.py
index e059ba011..b4741871e 100644
--- a/projects/mock_transformers/init_env.py
+++ b/projects/mock_transformers/init_env.py
@@ -111,16 +111,3 @@ def flow_softmax(*args, **kwargs):
 
 
 nn.functional.softmax = flow_softmax
-
-# -----------------mock flow.tensor---------------
-temp_tensor_func = flow.tensor
-
-
-def flow_tensor(input_x, **kwargs):
-    if isinstance(input_x, (int, float)):
-        return input_x
-    else:
-        return temp_tensor_func(input_x, **kwargs)
-
-
-flow.tensor = flow_tensor