   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);
+  prepare_training_data<<<1,1>>>(in_tensor, y_out, y_labels);
 
   for(int iter = 0; iter < 5000; iter++) {
 

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


  in_tensor->at(3, 0, 0, 0) = 1;

}


__global__ void prepare_training_data(Tensor<float> *in_tensor, Tensor<float> * y_out, Tensor<float> * y_label) {


  in_tensor->at(0, 0, 0, 0) = 0;
