Prefetch on GPU cleanup & Cleanup

Author: DEKHTIARJonathan

tftrt/examples/benchmark_runner.py

@@ -432,14 +432,14 @@ def log_step(step_idx, display_every, iter_time, memcpyHtoD_time, dequeue_time):
                                 print("[Exiting] Reached end of dataset ...")
                                 break
 
+                        with tracing_ctx('Inputs Preprocessing', step_num=step_idx, _r=1):
+                            x, y = self.preprocess_model_inputs(data_batch)
+
                         with tracing_ctx('Inputs MemcpyHtoD', step_num=step_idx, _r=1):
                             start_time = time.time()
-                            data_batch = force_data_on_gpu_fn(data_batch)
+                            x = force_data_on_gpu_fn(x)
                             memcopy_times.append(time.time() - start_time)
 
-                        with tracing_ctx('Inputs Preprocessing', step_num=step_idx, _r=1):
-                            x, y = self.preprocess_model_inputs(data_batch)
-
                         with tracing_ctx('GPU Inference', step_num=step_idx, _r=1):
                             start_time = time.time()
                             y_pred = infer_batch(x)

tftrt/examples/nvidia_examples/efficientnet_base/infer.py

@@ -221,13 +221,6 @@ def preprocess_model_inputs(self, data_batch):
 
         x, y = data_batch
 
-        x = {
-            "image": x['image'],
-            "mixup_weight": x['mixup_weight'],
-            "cutmix_mask": x['cutmix_mask'],
-            "is_tr_split": x['is_tr_split']
-        }
-
         return x, y
 
     def postprocess_model_outputs(self, predictions, expected):