 
     float mse = 0;
     for(int b = 0; b < batch_size; b++) {
-      mse += (out_tensor_diff->at(b,0,0,0)*out_tensor_diff->at(b,0,0,0) + out_tensor_diff->at(b,0,0,1)*out_tensor_diff->at(b,0,0,1) + out_tensor_diff->at(b,0,0,2)*out_tensor_diff->at(b,0,0,2))/3
+      mse += (out_tensor_diff->at(b,0,0,0)*out_tensor_diff->at(b,0,0,0) + out_tensor_diff->at(b,0,0,1)*out_tensor_diff->at(b,0,0,1) + out_tensor_diff->at(b,0,0,2)*out_tensor_diff->at(b,0,0,2))/3;
     }
     mse /= batch_size;
     printf("mse: %f \n", mse);
-    out.UpdateWb(0.02);
-    h1.UpdateWb(0.02);
+    out.UpdateWb(0.01);
+    h1.UpdateWb(0.01);
   }
 
 


    h1.Forward({in_tensor}, {h1_tensor});
    out.Forward({h1_tensor}, {out_tensor});

    for(int b = 0; b < batch_size; b++) {
      for(int i = 0; i < out_nodes; i++) {


  Tensor<float>* out_tensor = Tensor<float>::CreateTensorCPU(out_dims);
  Tensor<float>* out_tensor_diff = Tensor<float>::CreateTensorCPU(out_dims);
  Tensor<float>* softmax_out_tensor = Tensor<float>::CreateTensorCPU(softmax_out_dims);

  Tensor<float>* y_out = Tensor<float>::CreateTensorCPU(out_dims);

