diff --git a/test/test_fori_loop_simple_linear_model_test_code.py b/test/test_fori_loop_simple_linear_model_test_code.py
index b13e9f21abd..5e80deb2e67 100644
--- a/test/test_fori_loop_simple_linear_model_test_code.py
+++ b/test/test_fori_loop_simple_linear_model_test_code.py
@@ -92,8 +92,13 @@ def body_fun(l_in_i):
 # print("weight_: ", weight_)
 print("l_in_i_plus_1_: ", l_in_i_plus_1_)
 print("l_out_: ", l_out_)
-print("bias_: ", linear_0.bias)
-print("weight_: ", linear_0.weight)
+# print("bias_: ", linear_0.bias)
+# print("weight_: ", linear_0.weight)
+
+fake_l_in_0 = torch.ones([10], dtype=torch.float32, device=device)
+expectation = linear_0(fake_l_in_0)
+
+assert l_out_ == expectation
 
 # --- linear two ---
 # l_in_2 = torch.randn(10, device=xm.xla_device())