test

pytorch · Apr 10, 2024 · bfeeb62 · bfeeb62
1 parent b3235b2
commit bfeeb62
Showing 1 changed file with 7 additions and 2 deletions.
diff --git 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())