diff --git a/tb_plugins/profiling/tb_plugin/torch_tb_profiler/profiler/run_generator.py b/tb_plugins/profiling/tb_plugin/torch_tb_profiler/profiler/run_generator.py
index d2504d62fd11100a74a7af51af9ee99d835cf93b..d28ebf8b476d4de5f1557fdb59dbf74f95902339 100644
--- a/tb_plugins/profiling/tb_plugin/torch_tb_profiler/profiler/run_generator.py
+++ b/tb_plugins/profiling/tb_plugin/torch_tb_profiler/profiler/run_generator.py
@@ -344,29 +344,16 @@ class RunGenerator(object):
             reader = csv.DictReader(f)
             for row in reader:
                 data = []
-                data.append(row.get('Op Name'))
-                data.append(float(row.get('Task Duration(us)')))
+                data.append(row.get('Name'))
+                data.append(float(row.get('Duration(us)')))
                 pie['rows'].append(data)
         datas = {'total': pie}
         return datas
 
     def _generate_kernel_table(self):
-        display_columns = {
-            'Step ID': 'Step ID',
-            'Op Name': 'Name',
-            'Op Type': 'Type',
-            'Task Type': 'Accelerator Core',
-            'Task Start Time': 'Start Time',
-            'Task Duration(us)': 'Duration(us)',
-            'Task Wait Time(us)': 'Wait Time(us)',
-            'Block Dim': 'Block Dim',
-            'Input Shapes': 'Input Shapes',
-            'Input Data Types': 'Input Data Types',
-            'Input Formats': 'Input Formats',
-            'Output Shapes': 'Output Shapes',
-            'Output Data Types': 'Output Data Types',
-            'Output Formats': 'Output Formats'
-        }
+        display_columns = ('Step ID', 'Name', 'Type', 'Accelerator Core', 'Start Time', 'Duration(us)', 'Wait Time(us)',
+                           'Block Dim', 'Input Shapes', 'Input Data Types', 'Input Formats', 'Output Shapes',
+                           'Output Data Types', 'Output Formats')
         display_idxs = []
         table = {'columns': [], 'rows': []}
         result = {
@@ -378,16 +365,16 @@ class RunGenerator(object):
         path = self.profile_data.kernel_file_path
         datas = self._get_csv_data(path)
         for idx, column in enumerate(datas[0]):
-            if column == 'Op Name':
+            if column == 'Name':
                 self.name_idx = idx
-            elif column == 'Task Duration(us)':
+            elif column == 'Duration(us)':
                 self.duration_idx = idx
-            elif column == 'Task Type':
+            elif column == 'Type':
                 self.core_type_idx = idx
 
-            if display_columns.get(column) is not None:
+            if column in display_columns:
                 display_idxs.append(idx)
-                table['columns'].append({'type': 'string', 'name': display_columns[column]})
+                table['columns'].append({'type': 'string', 'name': column})
         table['rows'] = [self._handle_kernel_table_rows(display_idxs, ls) for idx, ls in
                          enumerate(datas) if idx != 0]
         return result