     cudaStatus = cudaGetLastError();
     checkCudaErrors(cudaStatus);
 
-    // show_tensor<<<1, 1>>>(conv1_top_diff);
+    conv1.UpdateWb(lr);
+    cudaStatus = cudaGetLastError();
+    checkCudaErrors(cudaStatus);
 
     cudaStatus = cudaDeviceSynchronize();
     checkCudaErrors(cudaStatus);
 
+    
   }
-
-
-
 
   printf("Prediction: \n");
   Tensor<float>* out = Tensor<float>::TensorGPUtoCPU(sm_top);

  checkCudaErrors(cudaStatus);


  printf("Loading weights ...\n");

  std::string model_path = "models/cifar10/model.txt";
  std::ifstream file(model_path);

  size_t conv1_w_dims[4] = {5,5,3,32};
  Tensor<float>* conv1_w = Tensor<float>::CreateTensorCPU(conv1_w_dims);
  load_to_conv<float>(conv1_w, file);
  Tensor<float>::DataArrayCPUtoGPU(conv1_w, conv1.W_);

  size_t conv1_b_dims[4] = {1,1,1,32};
  Tensor<float>* conv1_b = Tensor<float>::CreateTensorCPU(conv1_b_dims);
  load_to_bias<float>(conv1_b, file);
  Tensor<float>::DataArrayCPUtoGPU(conv1_b, conv1.b_);

  size_t conv2_w_dims[4] = {5,5,32,32};
  Tensor<float>* conv2_w = Tensor<float>::CreateTensorCPU(conv2_w_dims);
  load_to_conv<float>(conv2_w, file);
  Tensor<float>::DataArrayCPUtoGPU(conv2_w, conv2.W_);

  size_t conv2_b_dims[4] = {1,1,1,32};
  Tensor<float>* conv2_b = Tensor<float>::CreateTensorCPU(conv2_b_dims);
  load_to_bias<float>(conv2_b, file);
  Tensor<float>::DataArrayCPUtoGPU(conv2_b, conv2.b_);

   size_t conv3_w_dims[4] = {5,5,32,64};
  Tensor<float>* conv3_w = Tensor<float>::CreateTensorCPU(conv3_w_dims);
  load_to_conv<float>(conv3_w, file);
  Tensor<float>::DataArrayCPUtoGPU(conv3_w, conv3.W_);

  size_t conv3_b_dims[4] = {1,1,1,64};
  Tensor<float>* conv3_b = Tensor<float>::CreateTensorCPU(conv3_b_dims);
  load_to_bias<float>(conv3_b, file);
  Tensor<float>::DataArrayCPUtoGPU(conv3_b, conv3.b_);

  size_t fc4_w_dims[4] = {1,1,64,1024};
  Tensor<float>* fc4_w = Tensor<float>::CreateTensorCPU(fc4_w_dims);
  load_to_fc<float>(fc4_w, file);
  Tensor<float>::DataArrayCPUtoGPU(fc4_w, fc4.W_);

  size_t fc4_b_dims[4] = {1,1,1,64};
  Tensor<float>* fc4_b = Tensor<float>::CreateTensorCPU(fc4_b_dims);
  load_to_bias<float>(fc4_b, file);
  Tensor<float>::DataArrayCPUtoGPU(fc4_b, fc4.b_);

  size_t fc5_w_dims[4] = {1,1,10,64};
  Tensor<float>* fc5_w = Tensor<float>::CreateTensorCPU(fc5_w_dims);
  load_to_fc<float>(fc5_w, file);
  Tensor<float>::DataArrayCPUtoGPU(fc5_w, fc5.W_);

  size_t fc5_b_dims[4] = {1,1,1,10};
  Tensor<float>* fc5_b = Tensor<float>::CreateTensorCPU(fc5_b_dims);
  load_to_bias<float>(fc5_b, file);
  Tensor<float>::DataArrayCPUtoGPU(fc5_b, fc5.b_);


  printf("Forward inference .. \n");
  data_layer.Forward(std::vector<Tensor<float>*> (), data_tops);
  conv1.Forward({data_tops[0]}, {conv1_top});
  pool1.Forward({conv1_top}, {pool1_top});
  relu1.Forward({pool1_top}, {relu1_top});
  conv2.Forward({relu1_top}, {conv2_top});
  pool2.Forward({conv2_top}, {pool2_top});
  relu2.Forward({pool2_top}, {relu2_top});
  conv3.Forward({relu2_top}, {conv3_top});
  pool3.Forward({conv3_top}, {pool3_top});
  relu3.Forward({pool3_top}, {relu3_top});
  // flatten the tensor
  size_t relu3_top_dims_reshaped[4] = {relu3_top_dims[0], relu3_top_dims[3], relu3_top_dims[1], relu3_top_dims[2]};
  Tensor<float> * reshaped_relu3_top_cpu = Tensor<float>::CreateTensorCPU(relu3_top_dims_reshaped);
  Tensor<float> * relu3_top_cpu = Tensor<float>::TensorGPUtoCPU(relu3_top);
  for(int b = 0; b < relu3_top_dims_reshaped[0]; b++) {
    for(int c = 0; c < relu3_top_dims_reshaped[1]; c++) {
      for(int h = 0; h < relu3_top_dims_reshaped[2]; h++) {
        for(int w = 0; w < relu3_top_dims_reshaped[3]; w++) {
          reshaped_relu3_top_cpu->at(b, c, h, w) = relu3_top_cpu->at(b, h, w, c);
        }
      }
    }
  }
  Tensor<float> * reshaped_relu3_top = Tensor<float>::TensorCPUtoGPU(reshaped_relu3_top_cpu);
  Tensor<float> * reshaped_relu3_top_diff = Tensor<float>::CreateTensorGPU(to_fc4_dims);
  // flatten the tensor
  // Tensor<float>::ReshapeTensorGPU(relu3_top, to_fc4_dims);
  Tensor<float>::ReshapeTensorGPU(reshaped_relu3_top, to_fc4_dims);
  fc4.Forward({reshaped_relu3_top}, {fc4_top});
  // fc4.Forward({relu3_top}, {fc4_top});
  fc5.Forward({fc4_top}, {fc5_top});
  softmax.Forward({fc5_top}, {sm_top});
  cel_layer.Forward({sm_top, data_tops[1]}, {cel_top});


  for(int i = 0; i < 1; i++) {
    cel_layer.Backward({cel_top}, {cel_top}, {sm_top, data_tops[1]}, {sm_top_diff, cel_loss_diff});
    softmax.Backward({sm_top}, {sm_top_diff}, {fc5_top}, {fc5_top_diff});
    fc5.Backward({fc5_top}, {fc5_top_diff}, {fc4_top}, {fc4_top_diff});
    fc4.Backward({fc4_top}, {fc4_top_diff}, {reshaped_relu3_top}, {reshaped_relu3_top_diff});
    
    // fc4.Backward({fc4_top}, {fc4_top_diff}, {relu3_top}, {relu3_top_diff});
    // Tensor<float>::ReshapeTensorGPU(relu3_top, relu3_top_dims);
    // Tensor<float>::ReshapeTensorGPU(relu3_top_diff, relu3_top_dims);

    cudaStatus = cudaGetLastError();
    checkCudaErrors(cudaStatus);
    conv2.Backward({conv2_top}, {conv2_top_diff}, {relu1_top}, {relu1_top_diff});
    cudaStatus = cudaGetLastError();
    checkCudaErrors(cudaStatus);

    conv2.UpdateWb(lr);
    cudaStatus = cudaGetLastError();
    checkCudaErrors(cudaStatus);

    relu1.Backward({relu1_top}, {relu1_top_diff}, {pool1_top}, {pool1_top_diff});
    cudaStatus = cudaGetLastError();
    checkCudaErrors(cudaStatus);
