diff --git a/llvm/projects/hpvm-tensor-rt/dnn_sources/include/utils.h b/llvm/projects/hpvm-tensor-rt/dnn_sources/include/utils.h
index c7ab713823020c2fe8022839e91955d414457116..225b75a8573132e76f72e36fcd2bea7c71fe9d70 100644
--- a/llvm/projects/hpvm-tensor-rt/dnn_sources/include/utils.h
+++ b/llvm/projects/hpvm-tensor-rt/dnn_sources/include/utils.h
@@ -28,7 +28,6 @@ void printTensorInfo(void* tensor_ptr){
 void dumpWeightsToFile(char* file_name, void* weights_ptr){
 
   struct Tensor* weights = (Tensor*) weights_ptr;
-
   // Move data back to host
   hpvm_request_tensor(weights, 0);
   
@@ -237,12 +236,12 @@ struct Tensor* readTrainedWeights(char* file_name, int data_type,
   fseek(file, file_header_size, SEEK_CUR); // Skipping the file header
   size_t bytes_read = fread(tensor_data, 1, size_in_bytes, file);
 
-  for(unsigned i = 0 ; i < 40; i++){
-    printf("tensor_data[%d] = %f \n", i, tensor_data[i]);
-  }
+  //for(unsigned i = 0 ; i < 40; i++){
+  //  printf("tensor_data[%d] = %f \n", i, tensor_data[i]);
+  //}
   
   struct Tensor* weights = (struct Tensor*) create4DTensor(data_type, nchw, dim1_size, dim2_size,
-					  dim3_size, dim4_size);
+					                   dim3_size, dim4_size);
   
   initTensorData(weights, tensor_data, size_in_bytes);
   compareValues(weights, tensor_data, num_elems);
@@ -253,7 +252,6 @@ struct Tensor* readTrainedWeights(char* file_name, int data_type,
 
 uint8_t* readLabels(char* labels_file, int num_labels){
 
-  //-- int file_header_size = 8;
   uint8_t* labels = (uint8_t*) malloc(sizeof(uint8_t) * num_labels);
   FILE* file = fopen(labels_file, "rb");
   if(file == NULL){
@@ -261,7 +259,6 @@ uint8_t* readLabels(char* labels_file, int num_labels){
     abort();
   }
 
-  //-- fseek(file, file_header_size, SEEK_CUR); // Skipping the file header
   size_t bytes_read = fread(labels, 1, sizeof(uint8_t) * num_labels, file);
 
   for(unsigned int i = 0 ; i < 20; i++){
@@ -315,7 +312,7 @@ void computeAccuracy(char* labels_file, int num_labels, void* result_ptr){
 
 
 
-void computeAccuracy2(uint8_t* labels, int num_labels, void* result_ptr){
+void computeAccuracy2(uint8_t* labels, int num_labels, void* result_ptr, unsigned num_classes = 10){
 
   unsigned num_zeros = 0;
   
@@ -326,10 +323,12 @@ void computeAccuracy2(uint8_t* labels, int num_labels, void* result_ptr){
   size_t channels = result->dims.dim_sizes[1];
   float* data = (float*) result->host_data;
   int num_errors = 0;
+
+  printf("batch_dim = %d, channels = %d \n", batch_dim, channels);
   
   for(int i = 0; i < batch_dim; i++){
     int chosen = 0;
-    for (int id = 1; id < 10; ++id){
+    for (int id = 1; id < num_classes; ++id){
       if (data[i * channels + chosen] < data[i * channels + id]) chosen = id;
     }