   h1.GetTopsDims({in_dims}, {h1_dims});
   size_t out_dims[4];
   out.GetTopsDims({h1_dims}, {out_dims});
+  size_t label_dims[4];
+  label_dims[0] = out_dims[0]; label_dims[1] = out_dims[1]; label_dims[2] = out_dims[2]; label_dims[3] = 1; 
   size_t softmax_out_dims[4];
   softmax_layer.GetTopsDims({out_dims}, {softmax_out_dims});
   size_t ce_out_dims[4];

  y_out->at(3,0,0,0) = 1;
  y_out->at(3,0,0,1) = 0;
  y_out->at(3,0,0,2) = 0;
}

void test_fc_bp_gpu() {

  Tensor<float>* cel_loss_diff = Tensor<float>::CreateTensorGPU(ce_loss_dims);

  Tensor<float>* y_out = Tensor<float>::CreateTensorGPU(out_dims);
  Tensor<float>* y_labels = Tensor<float>::CreateTensorGPU(label_dims);

  std::vector<std::vector<float>> x_train = {{0,1},{0,0},{1,0},{1,1}};
  std::vector<std::vector<float>> y_train = {{0,1,0}, {0,0,1}, {0,1,0}, {1,0,0}};

  prepare_training_data<<<1,1>>>(in_tensor, y_out, y_labels);

  for(int iter = 0; iter < 5000; iter++) {

