   FC<float> h1(in_nodes, h1_nodes, &const_init);
   FC<float> out(h1_nodes, out_nodes, &const_init);
   Softmax<float> softmax_layer;
+  CrossEntropyLoss<float> cel_layer;
 
   size_t in_dims[4] = {batch_size, 1, 1, in_nodes};
   size_t h1_dims[4];

  cudaFree(out_tensor_diff);
  cudaFree(softmax_out_tensor);
  cudaFree(softmax_out_tensor_diff);
}




    h1.Forward({in_tensor}, {h1_tensor});
    out.Forward({h1_tensor}, {out_tensor});
    // softmax_layer.Forward({out_tensor}, {softmax_out_tensor});
    cel_layer.Forward({out_tensor}, {cel_out_tensor});

    cudaStatus = cudaGetLastError();
    checkCudaErrors(cudaStatus);

    // calc_out_diff<<<1,1>>>(out_tensor_diff, out_tensor, y_out);
    // calc_out_diff<<<1,1>>>(softmax_out_tensor_diff, softmax_out_tensor, y_out);
    calc_out_diff<<<1,1>>>(cel_out_tensor_diff, cel_out_tensor, y_out);
    cudaStatus = cudaGetLastError();
    checkCudaErrors(cudaStatus);

    // softmax_layer.Backward({softmax_out_tensor}, {softmax_out_tensor_diff}, {out_tensor}, {out_tensor_diff});
    cel_layer.Backward({cel_out_tensor}, {cel_out_tensor_diff}, {out_tensor}, {out_tensor_diff});
    cudaStatus = cudaGetLastError();
    checkCudaErrors(cudaStatus);

